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

Question

我有两个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数组，其元素是其本身(从0开始)和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.])