numpy,获得最大的子集

Question

我有值的阵列,所述v(例如v=[1,2,3,4,5,6,7,8,9,10])和索引的阵列,说g(例如 g=[0,0,0,0,1,1,1,1,2,2]).

我知道,例如,如何以非常简单的方式采取每个组的第一个元素:

import numpy as np
v=np.array([1,2,3,4,74,73,72,71,9,10])
g=np.array([0,0,0,0,1,1,1,1,2,2])
mask=np.concatenate(([True],np.diff(g)!=0))
v[mask]

收益:

array([1, 74, 9])

是否有任何numpythonic方式(避免显式循环)来获得每个子集的最大值？

---

测试:

因为我收到了两个很好的答案,一个是python map,一个是numpy例程,我正在搜索性能最好的,这里有一些时间测试:

import numpy as np
import time
N=10000000
v=np.arange(N)
Nelemes_per_group=10
Ngroups=N/Nelemes_per_group
s=np.arange(Ngroups)
g=np.repeat(s,Nelemes_per_group)

start1=time.time()
r=np.maximum.reduceat(v, np.unique(g, return_index=True)[1])
end1=time.time()
print('END first method, T=',(end1-start1),'s')

start3=time.time()
np.array(list(map(np.max,np.split(v,np.where(np.diff(g)!=0)[0]+1))))
end3=time.time()
print('END second method,  (map returns an iterable) T=',(end3-start3),'s')

结果我得到:

END first method, T= 1.6057236194610596 s
END second method,  (map returns an iterable) T= 8.346540689468384 s

有趣的是,该map方法的大部分减速都是由于list()呼叫造成的.如果我不尝试将我的map结果重新转换为list(但我必须,因为python3.x返回一个迭代器:https: //docs.python.org/3/library/functions.html#map)

Answer 1

您可以使用np.maximum.reduceat：

>>> _, idx = np.unique(g, return_index=True)
>>> np.maximum.reduceat(v, idx)
array([ 4, 74, 10])

reduceat可以在此处找到有关 ufunc方法工作原理的更多信息。

性能备注

np.maximum.reduceat非常快。生成索引idx是这里花费的大部分时间。

虽然_, idx = np.unique(g, return_index=True)是一种获取索引的优雅方式，但它并不是特别快。

原因是np.unique需要先对数组进行排序，复杂度为O(n log n)。对于大型数组，这比使用多个 O(n) 操作来生成idx.

因此，对于大型数组，使用以下代码要快得多：

idx = np.concatenate([[0], 1+np.diff(g).nonzero()[0]])
np.maximum.reduceat(v, idx)