numpy summing matrix-rows by indices

Question

I have 3 matrices (np arrays):
A is of shape (n,m) ; B is of shape (m,k); and C is of shape (n,k)

Matrix C has only values from the set {-1,0,1} and it's an "indicator" of some sort: if C[i,j]==1 then I want to add the i-th row of A to the j-th column of b; and if C[i,j]==(-1) then subtract (0 does nothing).

It can easily be done with loops, but I'm wondering if there's a vectorized way to do it that might be faster?

example code:

C = np.array([[-1,  0,  0,  0,  1],
              [ 0,  0,  0,  0, -1],
              [ 0,  0,  0,  0, -1],
              [-1,  0,  0,  1,  1]])
a,b = np.where(C==1)
#here a=[0,3,3] and b=[4,3,4]
A[a,:] = [[0, 1, 2, 3, 4, 5, 6],
          [3, 3, 3, 3, 3, 3, 3],
          [3, 3, 3, 3, 3, 3, 3]]
B[:,b] += A[a]  #B is all 0.0 before  

expected result:

array([[ 0.,  0.,  0.,  3.,  3.],
       [ 0.,  0.,  0.,  3.,  4.],
       [ 0.,  0.,  0.,  3.,  5.],
       [ 0.,  0.,  0.,  3.,  6.],
       [ 0.,  0.,  0.,  3.,  7.],
       [ 0.,  0.,  0.,  3.,  8.],
       [ 0.,  0.,  0.,  3.,  9.]])

actual result:

array([[ 0.,  0.,  0.,  3.,  3.],
       [ 0.,  0.,  0.,  3.,  3.],
       [ 0.,  0.,  0.,  3.,  3.],
       [ 0.,  0.,  0.,  3.,  3.],
       [ 0.,  0.,  0.,  3.,  3.],
       [ 0.,  0.,  0.,  3.,  3.],
       [ 0.,  0.,  0.,  3.,  3.]])

Answer 1

We can simply use np.add.at on the transposed view of B -

np.add.at(B.T, b, A[a])

Sample run -

In [39]: C = np.array([[-1,  0,  0,  0,  1],
    ...:               [ 0,  0,  0,  0, -1],
    ...:               [ 0,  0,  0,  0, -1],
    ...:               [-1,  0,  0,  1,  1]])
    ...: a,b = np.where(C==1)
    ...: A = np.zeros((4,7),dtype=int)
    ...: A[a,:] = np.array([[0, 1, 2, 3, 4, 5, 6],
    ...:           [3, 3, 3, 3, 3, 3, 3],
    ...:           [3, 3, 3, 3, 3, 3, 3]])

In [40]: # Initialize B
    ...: B = np.zeros((7,5),dtype=int)

In [41]: np.add.at(B.T, b, A[a])

In [42]: B
Out[42]: 
array([[0, 0, 0, 3, 3],
       [0, 0, 0, 3, 4],
       [0, 0, 0, 3, 5],
       [0, 0, 0, 3, 6],
       [0, 0, 0, 3, 7],
       [0, 0, 0, 3, 8],
       [0, 0, 0, 3, 9]])

As commented by @DSM, we can also use matrix-mutiplication and hence avoiding the step to get indices for C==1 -

A.T.dot(C==1)