最快的重复索引添加（np.add.at用于多个维度）

Question

假设我有一个a的数组(N,) N~O(1e9)。 然后我有一个形状b的结果数组(M1,M2,M2)。

然后我有一个形状f的索引数组(N,3)，其中每一行都是b的索引。

我想尽快完成的是以下结果，

for i in range(N):
    b[*f[i]] += a[i]

像b[f[:,0],f[:,1],f[:,2]] += a这样的花式索引不起作用，因为f中有重复的索引，因此它只会将每个索引的最后一次出现添加到b。

如果b是1D，那么我可以np.add.at(b,f,a)。

我正在寻找最快的选择。

Answer 1

要使用np.add.at，请传入一个索引元组，如此 -

np.add.at(b, (f[:,0], f[:,1], f[:,2]), a)

计算线性索引后，np.bincount的更快方法就是这样 -

lidx = np.ravel_multi_index(f.T,b.shape)
b += np.bincount(lidx, a, minlength=b.size).reshape(b.shape)