向另一个数组索引的数组添加向量化方法-Python / NumPy

Question

我有3个Numpy数组， a ， b和c 。 b和c是非常大的数组，并且具有相同的长度。 b每个元素为0、1或2， a的长度也为3。现在，我想知道是否有一种方法可以消除以下for循环：

for i in range(len(b)):
    a[b[i]] += c[i]

任何评论将不胜感激。

Answer 1

NumPy ufuncs at这种情况下有一个at方法：

numpy.add.at(a, b, c)

这样做是每个人都希望a[b] += c对索引数组b所做的事情，然后他们才看到它不起作用。

Answer 2

您可以将np.bincount用于此类基于ID的求和，如下所示-

a += np.bincount(b,c,minlength=a.size)

运行时测试-

In [136]: # Large arrays as inputs
     ...: a = np.random.rand(3)
     ...: c = np.random.rand(10000)
     ...: b = np.random.randint(0,3,10000)
     ...: 
     ...: # Make copies for testing
     ...: a1 = a.copy()
     ...: a2 = a.copy()
     ...: 

In [137]: def bincount_app(a, b, c): # bincount approach as func
     ...:     a += np.bincount(b,c,minlength=a.size)
     ...:     

In [138]: %timeit np.add.at(a1, b, c) # @user2357112's soln
1000 loops, best of 3: 1.29 ms per loop

In [139]: %timeit bincount_app(a2, b, c)
10000 loops, best of 3: 36.6 µs per loop