numpy,获取最大的子集

numpy, get maximum of subsets

我有一个值数组,例如 v,(例如 v=[1,2,3,4,5,6,7,8,9,10])和一个索引数组,例如 g(例如 g=[0,0,0,0,1,1,1,1,2,2])。

例如,我知道如何以非常 numpythonic 的方式获取每个组的第一个元素,这样做:

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]

returns:

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 returns 一个迭代器: https://docs.python.org/3/library/functions.html#map

您可以像下面那样创建您的面具并使用 map 功能:

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

如果你不想得到一个带有map的生成器,你可以使用列表推导在列表中获得相同的结果,并注意列表推导的迭代在解释器,如内置函数。

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

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

这是一种使用 masking and 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.])

您可以使用 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)