在numpy数组中连续添加值,不进行循环

Question

也许以前曾经问过,但我找不到它.有时我有索引I,我想相应地将这个索引连续添加到另一个数组的numpy数组中.例如:

A = np.array([1,2,3])
B = np.array([10,20,30])
I = np.array([0,1,1])
for i in range(len(I)):
    A[I[i]] += B[i]
print(A)

打印预期(正确)值:

[11 52  3]

而

A[I] += B
print(A)

导致预期的(错误的)答案

[11 32  3].

有没有办法以矢量化的方式做我想要的,没有循环？如果没有,这是最快的方法吗？

Answer 1

用途numpy.add.at:

>>> import numpy as np
>>> A = np.array([1,2,3])
>>> B = np.array([10,20,30])
>>> I = np.array([0,1,1])
>>> 
>>> np.add.at(A, I, B)
>>> A
array([11, 52,  3])

或者,np.bincount:

>>> A = np.array([1,2,3])
>>> B = np.array([10,20,30])
>>> I = np.array([0,1,1])
>>> 
>>> A += np.bincount(I, B, minlength=A.size).astype(int)
>>> A
array([11, 52,  3])

哪个更快？

要看.在这个具体的例子中add.at似乎稍微快一点,大概是因为我们需要在bincount解决方案中转换类型.

如果OTOH A并B为floatD类,然后bincount会更快.

Answer 2

你需要使用np.add.at:

A = np.array([1,2,3])
B = np.array([10,20,30])
I = np.array([0,1,1])

np.add.at(A, I, B)
print(A)

版画

array([11, 52, 3])

这在文档中注明:

ufunc.at(a,indices,b = None)

对'index'指定的元素在操作数'a'上执行无缓冲的就地操作.对于加法ufunc,此方法等效于[indices] + = b,除了为多次索引的元素累积结果.例如,[[0,0]] + = 1只会因缓冲而增加第一个元素,而add.at(a,[0,0],1)将增加第一个元素两次.