Search code examples
pythonnumpymatrixarray-broadcasting

Subtract a column vector from matrix at specified vector of columns using only broadcast


I want to subtract a column vector from a numpy matrix using another vector which is index of columns where the first column vector needs to be subtracted from the main matrix. For eg.

M = array([[  1,   2,   1,   1],
           [  2,   1,   1,   1],
           [  1,   1,   2,   1],
           [  2,   1,   1,   1],
           [  1,   1,   1,   2]])  # An example matrix

V = array([1, 1, 1, 1, 1]) # An example column vector

I = array([0, 3, 2, 3, 1, 3, 3]) # The index maxtrix

Now I want to subtract V from M at column numbers given in I. For eg. I[0] is 0, so subtract V from first column (zero index) of matrix M.

Similarly I[1] = 3, subtract V from fourth column (three index) of matrix M.

At the end of operation, since 3 occurs 4 times in I, so V will be subtracted from third column i.e. last column of M- 4 times.

I need to do this using only broadcast, no loops.

I have tried the following:

M[:, I] - V[np.newaxis, :].T

but it ends up broadcasting resultant matrix to have more columns than there are in M.


Solution

  • One can use bincount and outer

    >>> M - np.outer(V, np.bincount(I, None, M.shape[1]))
    array([[ 0,  1,  0, -3],
           [ 1,  0,  0, -3],
           [ 0,  0,  1, -3],
           [ 1,  0,  0, -3],
           [ 0,  0,  0, -2]])
    

    or subtract.at

    >>> out = M.copy()
    >>> np.subtract.at(out, (np.s_[:], I), V[:, None])
    >>> out
    array([[ 0,  1,  0, -3],
           [ 1,  0,  0, -3],
           [ 0,  0,  1, -3],
           [ 1,  0,  0, -3],
           [ 0,  0,  0, -2]])