numpy:用索引数组有效求和

Question

假设我有2个矩阵M和N(都有> 1列).我还有一个索引矩阵I,有2列 - 1表示M,1表示N.N的索引是唯一的,但M的索引可能出现不止一次.我想要执行的操作是,

for i,j in w:
  M[i] += N[j]

除了for循环之外,还有更有效的方法吗？

Answer 1

为了完整起见,在numpy的> = 1.8,你也可以使用np.add的at方法:

In [8]: m, n = np.random.rand(2, 10)

In [9]: m_idx, n_idx = np.random.randint(10, size=(2, 20))

In [10]: m0 = m.copy()

In [11]: np.add.at(m, m_idx, n[n_idx])

In [13]: m0 += np.bincount(m_idx, weights=n[n_idx], minlength=len(m))

In [14]: np.allclose(m, m0)
Out[14]: True

In [15]: %timeit np.add.at(m, m_idx, n[n_idx])
100000 loops, best of 3: 9.49 us per loop

In [16]: %timeit np.bincount(m_idx, weights=n[n_idx], minlength=len(m))
1000000 loops, best of 3: 1.54 us per loop

除了明显的性能劣势外,它还有几个优点:

1. np.bincount将其权重转换为双精度浮点数,.at将与您的数组的本机类型一起运行.这使得它成为处理例如复数的最简单选择.
2. np.bincount只有权重加在一起,你有一个at所有ufunc 的方法,所以你可以重复multiply,或logical_and,或任何你想要的.

但对于您的用例,np.bincount可能是要走的路.