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

Vectorize a conditional sum of elements in an array of indices and another array of elements to be added

我有 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.]

对其进行矢量化的正确方法是什么?谢谢。

我相信你可以用 bincount:

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

输出:

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