Получить суммы соответствующих строк с numpy

Question

Например, у меня есть:

a = np.array([0,0,1,1,1,2,3,3,3,3])
b = np.arange(20).reshape(10, 2)

Я хочу вычислить сумму строк в b в соответствии с соответствующими идентификаторами, определенными в a (только сумма строк с одинаковыми идентификаторами). Таким образом, вывод в этом случае будет c, который имеет форму 4x2.

c[0] = b[0] + b[1]
c[1] = b[2] + b[3] + b[4]
c[2] = b[5]
c[3] = b[6] + b[7] + b[8] + b[9]

Как я могу добиться этого без для l oop? Есть ли у numpy какой-нибудь причудливый метод индексации, чтобы добиться этого более эффективно? Спасибо.

Answer 1

Вы можете использовать np.add.at:

>>> out=np.zeros((4,2),int)
>>> np.add.at(out,a,b)
>>> out
array([[ 2,  4],
       [18, 21],
       [10, 11],
       [60, 64]])

Answer 2

Вы можете использовать маску для исключения / включения значений, которые будут добавлены:

import numpy as np
a = np.array([0,0,1,1,1,2,3,3,3,3])
b = np.arange(20).reshape(10, 2)

mask   = a[:,None]==np.arange(4)
result = np.sum(b[:,:,None]*mask[:,None,:],axis=0)

вывод:

print(mask*1)  # rows of b that will be added in each column 

[[1 0 0 0]
 [1 0 0 0]
 [0 1 0 0]
 [0 1 0 0]
 [0 1 0 0]
 [0 0 1 0]
 [0 0 0 1]
 [0 0 0 1]
 [0 0 0 1]
 [0 0 0 1]]

print(result)  # result: 4 sums for each column of b 

[[ 2 18 10 60]
 [ 4 21 11 64]]

print(result.transpose()) # vertical representation:

[[ 2  4]
 [18 21]
 [10 11]
 [60 64]]

Answer 3

Пожалуйста, попробуйте это.

a = np.array([0,0,1,1,1,2,3,3,3,3])
b = np.arange(20).reshape(10, 2)
e = np.split(b, np.cumsum(np.unique(a, return_counts=True)[1])[:-1])
c = np.array(list(map(sum, e)))