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])

有关 ufunc reduceat 方法的工作原理的更多信息，请参见 here。

---

关于性能的评论

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)

Answer 2

这是一种使用masking 和broadcasting 的卷积向量化方法，将每个组放入常规二维数组的行中，然后沿每行找到最大值 -

# Mask of valid numbers from each group to be put in a regular 2D array
counts = np.bincount(g)
mask = np.arange(counts.max()) < counts[:,None]

# Group each group into rows of a 2D array and find max along ech row
grouped_2Darray = np.empty(mask.shape)
grouped_2Darray.fill(np.nan)
grouped_2Darray[mask] = v
out = np.nanmax(grouped_2Darray,1)

示例运行 -

In [52]: g
Out[52]: array([0, 0, 0, 0, 1, 1, 1, 1, 2, 2])

In [53]: v
Out[53]: array([ 1,  2,  3,  4, 74, 73, 72, 71,  9, 10])

In [54]: grouped_2Darray # Notice how elements from v are stacked
Out[54]: 
array([[  1.,   2.,   3.,   4.],
       [ 74.,  73.,  72.,  71.],
       [  9.,  10.,  nan,  nan]])

In [55]: np.nanmax(grouped_2Darray,1)
Out[55]: array([  4.,  74.,  10.])

Answer 3

你可以像下面这样创建你的面具并使用map函数：

>>> mask=np.diff(g)!=0
>>> map(np.max,np.split(v,np.where(mask)[0]+1))
[4, 74, 10]

如果您不想使用 map 获得生成器，您可以使用列表推导在列表中获得相同的结果，并注意列表推导的迭代在解释器中以 C 语言速度执行，如内置-in 函数。

[np.max(arr) for arr in np.split(v,np.where(mask)[0]+1)]

但我认为 numpythonic 解决方案仍然更好用。