Суммируйте 2darray по индексу

Question

Я хочу суммировать столбцы двумерного массива dat по индексу строки idx.Следующий пример работает, но медленно для больших массивов.Любая идея, чтобы ускорить его?

import numpy as np

dat = np.arange(18).reshape(6, 3, order = 'F')
idx = np.array([0, 1, 1, 1, 2, 2])

for i in np.unique(idx):
    print(np.sum(dat[idx==i], axis = 0))

Вывод

[ 0  6 12]
[ 6 24 42]
[ 9 21 33]

Answer 1

Подход № 1

Мы можем использовать matrix-multiplication с np.dot -

In [56]: mask = idx[:,None] == np.unique(idx)

In [57]: mask.T.dot(dat)
Out[57]: 
array([[ 0,  6, 12],
       [ 6, 24, 42],
       [ 9, 21, 33]])

Подход № 2

Для случая с idx уже отсортированным, мы можем использовать np.add.reduceat -

In [52]: p = np.flatnonzero(np.r_[True,idx[:-1] != idx[1:]])

In [53]: np.add.reduceat(dat, p, axis=0)
Out[53]: 
array([[ 0,  6, 12],
       [ 6, 24, 42],
       [ 9, 21, 33]])

Answer 2

Немного более быстрый подход с set объектом и ndarray.sum() методом:

In [216]: for i in set(idx):
     ...:     print(dat[idx == i].sum(axis=0))
     ...:     
[ 0  6 12]
[ 6 24 42]
[ 9 21 33]

---

Сравнение времени выполнения:

In [217]: %timeit for i in np.unique(idx): r = np.sum(dat[idx==i], axis = 0)
109 µs ± 1.1 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

In [218]: %timeit for i in set(idx): r = dat[idx == i].sum(axis=0)
71.1 µs ± 1.98 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)