Сумма строк на основе индекса с Numpy

Question

У меня есть 2D-массив, состоящий из 2D-векторов и 1D-массива индексов.

Как добавить / сложить строки 2D-массива , которые имеют один и тот же индекс , используя numpy?

Пример:

arr = np.array([[48, -51], [-15, -55], [26, -49], [-13, -17], [-67,  -7], [23, -48], [-29, -64], [37,  68]])
idx = np.array([0, 1, 1, 2, 2, 3, 3, 4])

#desired output
array([[48, -51],
       [11, -104],
       [-80, -24],
       [-6, -112],
       [ 37, 68]])

Обратите внимание, что исходный массив arr имеет форму (8, 2), а результат операции (5, 2).

Answer 1

Если индексы не всегда сгруппированы, сначала примените np.argsort:

order = np.argsort(idx)

Расположение сумм можно вычислить с помощью np.diffзатем np.flatnonzero, чтобы получить индексы. Мы также добавим ноль и сместим все на 1:

breaks = np.flatnonzero(np.concatenate(([1], np.diff(idx[order])))

breaks теперь можно использовать в качестве аргумента для np.add.reduceat:

result = np.add.reduceat(arr[order, :], breaks, axis=0)

Если индексы уже сгруппированы, вам вообще не нужно использовать order:

breaks = np.flatnonzero(np.concatenate(([1], np.diff(idx)))
result = np.add.reduceat(arr, breaks, axis=0)

Answer 2

Вы можете использовать pandas для цели:

pd.DataFrame(arr).groupby(idx).sum().to_numpy()

Вывод:

array([[  48,  -51],
       [  11, -104],
       [ -80,  -24],
       [  -6, -112],
       [  37,   68]])