向量化索引数组和要添加的另一个元素数组中元素的条件总和

Question

我有 2 个 numpy 数组，可能如下所示：n = np.array([[1, 2, 2], [4, 5, 5], [7, 7, 7]]) 和 F = np.array([[5, 6, 7], [8, 9, 1], [2, 3, 4]])，其中 n 是索引矩阵，F 是要添加到结果数组的元素或值的矩阵。 result 将是一个 1D numpy 数组，其元素是自身求和的结果（从零开始）以及由n 中的索引和F 中的值指示的相应值。换句话说，我想对以下操作进行矢量化：

n = np.array([[1, 2, 2], [4, 5, 5], [7, 7, 7]])
F = np.array([[5, 6, 7], [8, 9, 1], [2, 3, 4]])

nmax = n.max()

result = np.zeros((nmax + 1,))
ravel_n = n.ravel()
ravel_F = F.ravel()
for i in range(n.shape[0] * n.shape[1]):
    if ravel_n[i] > 0:
        result[ravel_n[i]] = result[ravel_n[i]] + ravel_F[i]

这给了我

>>> result
[ 0.  5. 13.  0.  8. 10.  0.  9.]

我尝试了以下矢量化表达式，但它输出的结果不正确。

>>> result_2 = np.zeros((nmax + 1,))
>>> result_2[ravel_n[ravel_n > 0]] = result_2[ravel_n[ravel_n > 0]] + ravel_F[ravel_n > 0]
>>> result_2
[0. 5. 7. 0. 8. 1. 0. 4.]

矢量化的正确方法是什么？谢谢。

Answer 1

我相信你可以用bincount：

np.bincount(n.ravel(), weights=F.ravel())

输出：

array([ 0.,  5., 13.,  0.,  8., 10.,  0.,  9.])