numpy, получить максимум подмножеств

у меня есть массив значений, сказал 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])

есть numpyтонический способ (избегая явных циклов), чтобы получить максимум каждого подмножества?


тесты:

так как я получил два хороших ответы, один с питоном 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 )

3 ответов


можно использовать np.maximum.reduceat:

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

подробнее о работе ufunc reduceat способ можно найти здесь.


замечание о производительности

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)

вот один извилистый векторизованный подход с использованием masking и broadcasting это помещает каждую группу в строки регулярного 2D-массива, а затем находит максимум вдоль каждой строки -

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

вы можете создать маску, как следующие и использовать map функция :

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

если вы не хотите получить генератор с map вы можете использовать понимание списка для достижения того же результата в списке, и обратите внимание, что итерация понимания списка выполнялась со скоростью языка C внутри интерпретатора, как встроенные функции.

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

но я думаю, что numpythonic решение по-прежнему лучше использовать.