索引数组中重复条目的总和

时间:2020-03-02 01:11:04

标签: python arrays numpy indexing sum

给出numpy ndarray A和一个整数数组I,它们的形状相同,具有最高值imax和数组B = np.zeros(imax),我们可以做B[I] = A 。但是,如果I已重复输入,则最后的分配成立。我需要在对重复的条目求和时进行此操作,例如

For i in range(A.size):
    B[I.ravel()[i]] += A.ravel()[i]

numpy中是否有一种很好的方法?

例如,我想要这种行为(但是=+=都没有这样的作用

A = np.array((1,2,5,9))
I = np.array((0,1,2,0),dtype=int)
B = np.zeros(3)
B[I] += A
print(B)
>>> array([10,2,5])

在这里,我们在第一项中看到1+9=10

1 个答案:

答案 0 :(得分:2)

In [1]: A = np.array((1,2,5,9)) 
   ...: I = np.array((0,1,2,0),dtype=int) 
   ...: B = np.zeros(3) 
   ...: B[I] += A                                                                                                          
In [2]: B                                                                                                                  
Out[2]: array([9., 2., 5.])

这是一种缓冲解决方案,与迭代解决方案不同:

In [3]: B = np.zeros(3)                                                                                                    
In [4]: for i,a in zip(I,A): 
   ...:     B[i] += a 
   ...:                                                                                                                    
In [5]: B                                                                                                                  
Out[5]: array([10.,  2.,  5.])

使用ufunc.at的无缓冲解决方案:

In [6]: B = np.zeros(3)                                                                                                    
In [7]: np.add.at(B, I, A)                                                                                                 
In [8]: B                                                                                                                  
Out[8]: array([10.,  2.,  5.])