Блог

Каков наилучший способ векторизации следующего кода в numpy?

 from numpy import *

A = zeros(5, dtype='int')
I = [1, 1, 1, 3]
J = [2, 1, 1, 1]
for i, j in zip(I, J):
    A[i]  = j

print A
  

Результат должен быть:

 [0 4 0 1 0]
  

Здесь A находится исходный массив, I в котором хранится индекс, на который мы хотим увеличить соответствующую запись J .

Если просто векторизировать вышесказанное, выполнив:

 A[I]  = J
print A
  

один из них дает неправильный ответ

 [0 1 0 1 0]
  

поскольку, по-видимому, повторяющиеся индексы игнорируются. Существует ли эквивалентная операция = , которая не игнорирует повторяющиеся индексы?

Вы можете использовать numpy.bincount() :

 A = numpy.zeros(5, dtype='int')
I = [1, 1, 1, 3]
J = [2, 1, 1, 1]
sums = numpy.bincount(I, J)
A[:len(sums)]  = sums
print(A)
  

С принтами

 [0 4 0 1 0]
  

В принципе, вы можете сделать это с помощью numpy ‘s bincount и unique , но я бы предположил, что это только сделает код намного менее читаемым без какого-либо ощутимого улучшения производительности.