Векторизация операции индексации в numpy

Question

Каков наилучший способ векторизации следующего кода в 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]

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

Answer 1

Вы можете использовать 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]

Answer 2

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