NumPy：使用索引数组高效求和

Question

NumPy：使用索引数组高效求和

11

假设我有两个矩阵M和N（它们都有>1列）。我还有一个索引矩阵I，其中包括2列--一列是针对M的，另一列是针对N的。 N的索引是唯一的，但M的索引可能会出现多次。我想执行的操作是

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

除了使用for循环，是否有更有效的方式来完成这个任务？

- duckworthd

问题中缺少的是：显然的M[w[:, 0]] += N[w[:, 1]]会失败，因为w[:, 0]中的索引不是唯一的，所以值会被覆盖。 - Fred Foo

当您不进行原地加法而创建第三个数组时，会发生什么？这样做是否会给您带来所需的结果？ - Midnighter

1

如果您不在原地执行操作，则最终会得到一个新的数组，其长度与I匹配，其中I[i] = M[w[i,0]] + N[w[i,1]]。那么，如何按w[:,0]分组并对其值求和？ - duckworthd

3个回答

3

同时使用 m_ind，n_ind = w.T，然后只需执行 M += np.bincount(m_ind, weights=N[n_ind], minlength=len(M))

- seberg

比我的解决方案更优雅；我有一种直觉，bincount 可以解决这个问题，但我无法找到规律。 - Fred Foo

我认为只有当N是一维数组时才能起作用。如果N是矩阵，np.bincount会抛出ValueError异常。 - duckworthd

1

为了清晰起见，让我们定义：

>>> 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])

- Fred Foo

@Veedrac 再试一次，它确实有效。

m_ind, n_ind = w.T; N = N[n_ind]; update = csr_matrix((N, m_ind, [0, len(N)])).toarray(); M[np.unique(m_ind)] += update.ravel()

。 - Fred Foo

这让我从M[np.unique(m_ind)] += update部分得到了ValueError: non-broadcastable output operand with shape (1,) doesn't match the broadcast shape (1,1)。只是错过了ravel，没关系。 - Veedrac

@Veedrac：抱歉，关于这个问题我已经在回答中解释了。 - Fred Foo

网页内容由stack overflow 提供, 点击上面的

可以查看英文原文，
原文链接

- Jaime · Accepted Answer

为了完整性，在 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

除了明显的性能劣势外，它还有一些优点：

np.bincount 将其权重转换为双精度浮点数，而 .at 会使用数组的本机类型进行操作。这使它成为处理复杂数字等问题的最简单选项。
np.bincount 仅将权重相加，您可以使用所有 ufunc 的 at 方法，因此您可以反复执行 multiply、logical_and 或其他您想要的操作。

但对于您的用例，np.bincount 可能是最佳选择。