Search code examples
numpymatrix-indexingnumpy-indexing

Assign zero to multiple columns in each row in a 2D numpy array efficiently

I want to assign zero to a 2d numpy matrix where for each row we have a cutoff column after which should be set to zero.

For example here we have matrix A of size 4x5 with cutoff columns [1,3,2,4]. What I want to do:

import numpy as np
np.random.seed(1)
A = np.random.rand(4, 5)
cutoff = np.array([1,3,2,4])
A[0, cutoff[0]:] = 0
A[1, cutoff[1]:] = 0
A[2, cutoff[2]:] = 0
A[3, cutoff[3]:] = 0

I could do it with np.repeat ing row indices and columns, but my matrix is so big I can't do that. Is there an efficient way to do this?

Solution

  • Use broadcasting to create the full mask and assign -

    A[cutoff[:,None] <= np.arange(A.shape[1])] = 0

    Alternatively with builtin outer method -

    A[np.less_equal.outer(cutoff, range(A.shape[1]))] = 0