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方法：

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

添加到由另一个数组索引的数组的矢量化方式 - Python / NumPy

2 个答案: