add row from matrix a to matrix b by index in python

Question

I am trying to do the following computation in a more efficient (vectorized) way instead of a loop.

for i in range(N):

  y[:,idx[i]] -= X[i]

where:

• X is N by D matrix,
• y is k by D matrix
• idx is a vector of length N with values between 0 and k-1

Example:

x = array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11],
       [12, 13, 14, 15],
       [16, 17, 18, 19]])


y =  array([[100, 101, 102, 103],
           [104, 105, 106, 107],
           [108, 109, 110, 111]])

idx = array([2, 1, 1, 1, 2])

So the output should be:

y = array([[100, 101, 102, 103],
          [ 80,  78,  76,  74],
          [ 92,  91,  90,  89]])

and this is equivalent to:

>>> for i in range(5):
...     y[idx[i],:] -= x[i]

Answer 1

Pandas solution is the shortest to go. You can do it via NumPy but a bit longer code:

import pandas as pd
df = pd.DataFrame(x, idx)
s = df.groupby(df.index).sum()
y[s.index] -= s

or in your case, since index are indices, you can directly use:

df = pd.DataFrame(x)
s = df.groupby(idx).sum()
y[s.index] -= s

output:

array([[100, 101, 102, 103],
       [ 80,  78,  76,  74],
       [ 92,  91,  90,  89]])