A pythonic way to mask all values after specific index in a 2D numpy array with 0

Question

I have a 2D array and a 1D index array like following:

a = np.random.rand(2,5)
a
>>> a
    array([[0.70095892, 0.01068342, 0.69875872, 0.95125273, 0.18589609],
   [0.02990893, 0.78341353, 0.12445391, 0.71709479, 0.24082166]])
>>> ind
    array([3, 2])

I want all values in row 1 after (including) index 3 become 0 and all values in row 2 after index 2 become 0. So the final output be:

array([[0.70095892, 0.01068342, 0.69875872, 0, 0],
   [0.02990893, 0.78341353, 0, 0, 0]])

Can you help me wih this?

Answer 1

You can do this via boolean indexing:

import numpy as np
a = np.random.rand(2, 5)
ind = np.array([3, 2])

# create a boolean indexing matrix
bool_ind = np.arange(a.shape[1]) >= ind[:, None]
print(bool_ind)
# [[False False False  True  True]
#  [False False  True  True  True]]

# modify values in a using the boolean matrix
a[bool_ind] = 0
print(a)
# [[0.78594869 0.11185728 0.06070476 0.         0.        ]
#  [0.48258651 0.3223349  0.         0.         0.        ]]