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 只是将权重加在一起，你有一个适用于所有 ufunc 的 at 方法，所以你可以重复 multiply，或 logical_and，或任何你喜欢的。

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

Answer 2

也使用m_ind, n_ind = w.T，只需使用M += np.bincount(m_ind, weights=N[n_ind], minlength=len(M))

Answer 3

为了清楚起见，让我们定义

>>> m_ind, n_ind = w.T

然后是for 循环

for i, j in zip(m_ind, n_ind):
    M[i] += N[j]

更新条目M[np.unique(m_ind)]。写入其中的值为N[n_ind]，必须按m_ind 分组。 （除了m_ind 之外还有一个n_ind 的事实实际上与问题无关；您可以设置N = N[n_ind]。）恰好有一个SciPy 类可以做到这一点：scipy.sparse.csr_matrix。

示例数据：

>>> m_ind, n_ind = array([[0, 0, 1, 1], [2, 3, 0, 1]])
>>> M = np.arange(2, 6)
>>> N = np.logspace(2, 5, 4)

for 循环的结果是 M 变为 [110002 1103 4 5]。我们使用csr_matrix 得到相同的结果，如下所示。正如我之前所说，n_ind 无关紧要，所以我们先去掉它。

>>> N = N[n_ind]
>>> from scipy.sparse import csr_matrix
>>> update = csr_matrix((N, m_ind, [0, len(N)])).toarray()

CSR 构造函数在所需索引处构建具有所需值的矩阵；其参数的第三部分是压缩列索引，这意味着值N[0:len(N)] 具有索引m_ind[0:len(N)]。重复项相加：

>>> update
array([[ 110000.,    1100.]])

这个形状为(1, len(np.unique(m_ind)))，可以直接添加进去：

>>> M[np.unique(m_ind)] += update.ravel()
>>> M
array([110002,   1103,      4,      5])