Question

我想通过行索引dat对2d数组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

我们可以利用np.dot和matrix-multiplication-

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)

按索引汇总2darray

2 个答案: