numpy where condition on index

November 6, 2021

I have a numpy 1D array of numbers representing columns e.g: [0,0,2,1]
And a matrix e.g:

[[1,1,1],
[1,1,1],
[1,1,1],
[1,1,1]]

Now I a want to change the values in the matrix to 0 where the column index is bigger than the value given in the 1D array:

[[1,0,0],
[1,0,0],
[1,1,1],
[1,1,0]]

How can I achieve this? I think I need a condition based on the index, not on the value

Explanation:
The first row in the matrix has indices [0,0 ; 0,1 ; 0,2] where the second index is the column. For indices 0,0 ; 0,1 and 0,2 the value 0 is given. 1 and 2 are bigger than 0. Thus only 0,0 is not changed to zero.

>Solution :

Assuming a the 2D array and v the 1D vector, you can create a mask of the same size and use numpy.where:

x,y = a.shape
np.where(np.tile(np.arange(y), (x,1)) <= v[:,None], a, 0)

Input:

a = np.array([[1,1,1],
              [1,1,1],
              [1,1,1],
              [1,1,1]])

v = np.array([0,0,2,1])

Output:

array([[1, 0, 0],
       [1, 0, 0],
       [1, 1, 1],
       [1, 1, 0]])

Intermediates:

>>> np.tile(np.arange(y), (x,1))
[[0 1 2]
 [0 1 2]
 [0 1 2]
 [0 1 2]]

>>> np.tile(np.arange(y), (x,1)) <= v[:,None]
[[ True False False]
 [ True False False]
 [ True  True  True]
 [ True  True False]]