Python numpy：从第二个数组中选择性地给定信息对数组求和

Question

Let's say I have a N-dimensionnal array, for example:

A = [ [ 1, 2] ,
    [6, 10] ]

and another array B that defines an index associated with each value of A

B = [[0, 1], [1, 0]]

And I want to obtain a 1D list or array that for each index contains the sum of the values of A associated with that index. For our example, we would want

C = [11, 8]

Is there a way to do this efficiently, without looping over the arrays manually ?

Edit: To make it clearer what I want, if we now take A the same and B equal to :

B = [[1, 1], [1,1]]

Then I want all the values of A to sum into the index 1 of C, which yields

C = [0, 19]

Or I can write a code snippet :

C = np.zeros(np.max(B))
for i in range(...):
   for j in range(...):
      C[B[i,j]] += A[i,j]
return C

Answer 1

For each index, or rather tag , you first get a True / False array of B where it's equal to the tag. Then you use numpy.nonzero to get indices by dimension of where the True s are in that array. When you index A with these indices, you get a new array with only the tagged elements of A . Finally you can sum over it.

[np.sum(A[np.nonzero(B==tag)]) for tag in [0,1]]

You could have also just done np.sum(A*(B==tag)) because the False s multiply as 0s, but then you'd be summing over a lot of 0s.

Answer 2

I think I found the best answer for now actually. I can just use np.histogram(B, weights = A) which should do the operation I want.