Предположим, что у меня есть два массива
A = [ 6, 4, 5, 7, 9 ]
ind = [ 0, 0, 2, 1, 2 ]
и функции f.
Я хочу построить новый массив B размера числа различных элементов в ind с B [i] результатом f с параметром является подмассиво A с индексом i.
В этом примере, если взять f = sum, то
B = [10, 7, 14]
или f = max
B = [6, 7, 9]
Есть ли более эффективный способ, чем цикл for в numpy?
Спасибо
Для частного случая f = sum
:
In [32]: np.bincount(ind,A)
Out[32]: array([ 10., 7., 14.])
Предполагая, что:
f
- это ufunclen(A) x len(A)
Вы можете сделать 2D-массив B
:
B=np.zeros((len(A),max(ind)+1))
и заполните в разных местах B
значения из A
, так что первый столбец B
получает значения A
только при ind == 0
, а второй столбец B
только получает значения от A
, когда ind == 1
и т.д.:
B[zip(*enumerate(ind))]=A
у вас получится массив вроде
[[ 6. 0. 0.]
[ 4. 0. 0.]
[ 0. 0. 5.]
[ 0. 7. 0.]
[ 0. 0. 9.]]
Затем вы можете применить f
вдоль axis=0, чтобы получить желаемый результат.
Здесь используется третье предположение:
B
не влияют желаемый результат.Если вы можете перенести эти предположения, то:
import numpy as np
A = np.array([ 6, 4, 5, 7, 9 ])
ind = np.array([ 0, 0, 2, 1, 2 ])
N=100
M=10
A2 = np.array([np.random.randint(M) for i in range(N)])
ind2 = np.array([np.random.randint(M) for i in range(N)])
def use_extra_axis(A,ind,f):
B=np.zeros((len(A),max(ind)+1))
B[zip(*enumerate(ind))]=A
return f(B)
def use_loop(A,ind,f):
n=max(ind)+1
B=np.empty(n)
for i in range(n):
B[i]=f(A[ind==i])
return B
def fmax(arr):
return np.max(arr,axis=0)
if __name__=='__main__':
print(use_extra_axis(A,ind,fmax))
print(use_loop(A,ind,fmax))
При определенных значениях M
и N
(например, M=10, N=100) использование дополнительной оси может быть быстрее, чем использование цикла:
% python -mtimeit -s'import test,numpy' 'test.use_extra_axis(test.A2,test.ind2,test.fmax)'
10000 loops, best of 3: 162 usec per loop
% python -mtimeit -s'import test,numpy' 'test.use_loop(test.A2,test.ind2,test.fmax)'
1000 loops, best of 3: 222 usec per loop
Однако, когда N растет больше (скажем M=10, N=100 00), использование цикла может быть быстрее:
% python -mtimeit -s'import test,numpy' 'test.use_extra_axis(test.A2,test.ind2,test.fmax)'
100 loops, best of 3: 13.9 msec per loop
% python -mtimeit -s'import test,numpy' 'test.use_loop(test.A2,test.ind2,test.fmax)'
100 loops, best of 3: 4.4 msec per loop
Включение отличная идея в использовании разреженной матрицы:
def use_sparse_extra_axis(A,ind,f):
B=scipy.sparse.coo_matrix((A, (range(len(A)), ind))).toarray()
return f(B)
def use_sparse(A,ind,f):
return [f(v) for v in scipy.sparse.coo_matrix((A, (ind, range(len(A))))).tolil().data]
Какая реализация лучше всего зависит от параметров N
и M
:
N=1000, M=100
·───────────────────────·────────────────────·
│ use_sparse_extra_axis │ 1.15 msec per loop │
│ use_extra_axis │ 2.79 msec per loop │
│ use_loop │ 3.47 msec per loop │
│ use_sparse │ 5.25 msec per loop │
·───────────────────────·────────────────────·
N=100000, M=10
·───────────────────────·────────────────────·
│ use_sparse_extra_axis │ 35.6 msec per loop │
│ use_loop │ 43.3 msec per loop │
│ use_sparse │ 91.5 msec per loop │
│ use_extra_axis │ 150 msec per loop │
·───────────────────────·────────────────────·
N=100000, M=50
·───────────────────────·────────────────────·
│ use_sparse │ 94.1 msec per loop │
│ use_loop │ 107 msec per loop │
│ use_sparse_extra_axis │ 170 msec per loop │
│ use_extra_axis │ 272 msec per loop │
·───────────────────────·────────────────────·
N=10000, M=50
·───────────────────────·────────────────────·
│ use_loop │ 10.9 msec per loop │
│ use_sparse │ 11.7 msec per loop │
│ use_sparse_extra_axis │ 15.1 msec per loop │
│ use_extra_axis │ 25.4 msec per loop │
·───────────────────────·────────────────────·
Я не думаю, что вы можете уйти от цикла, но, возможно, использование scipy разреженных матриц будет более эффективным.
[f(v) for v in scipy.sparse.coo_matrix((A, (ind, range(len(A))))).tolil().data]
По крайней мере для добавления, это работает
import numpy as np
def op_at(f, ind, vals):
base = np.zeros(np.max(ind)+1)
f.at(base, ind, vals)
return base
print op_at(np.add, [ 0, 0, 2, 1, 2], [ 6, 4, 5, 7, 9])
> [ 10. 7. 14.]
К сожалению, он не работает для max.
Другая возможность,
from operator import itemgetter
from itertools import groupby
A = [ 6, 4, 5, 7, 9 ]
ind = [ 0, 0, 2, 1, 2 ]
z = zip(ind,A)
z.sort()
fst,snd = itemgetter(0), itemgetter(1)
g = groupby(z,fst)
f = sum
# or
# f = max
for i in g:
print i[0],f(snd(j) for j in i[1])