numpy where condition on index-CodePudding

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.

CodePudding user response：

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]]

CodePudding user response：

Construct a 2D array whose elements are the corresponding column index, and then mask the elements greater than the corresponding value of the 1D array.

Taking advantage of broadcasting, you can do:

>>> arr = np.ones((4,3))
>>> arr 

array([[1., 1., 1.],
       [1., 1., 1.],
       [1., 1., 1.],
       [1., 1., 1.]])

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

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

>>> mask = np.arange(arr.shape[1])[None,:] > col_thr_idx[:,None]
>>> mask

array([[False,  True,  True],
       [False,  True,  True],
       [False, False, False],
       [False, False,  True]])

>>> arr[mask] = 0
>>> arr

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