Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Numpy sum over repeated entries in index array

Given numpy ndarray A and an integer array I, of the same shape, with highest value imax and an array B = np.zeros(imax) we can do B[I] = A. However if I has repeated entries the last assignment holds. I need to do this while summing over repeated entries instead, like

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

Is there a good way to do this in numpy?

For example, I want this behavior (but neither = nor += works like this)

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])

Here we see 1+9=10 in the first entry.

like image 427
Jonatan Öström Avatar asked Sep 15 '25 19:09

Jonatan Öström


1 Answers

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.])

This a buffered solution, different from an iterative one:

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.])

The unbuffered solution using the 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.])
like image 108
hpaulj Avatar answered Sep 17 '25 08:09

hpaulj