numpy array mask with range of elements specified by another array row-wise

Question

Having a 2d numpy array, for example

data = np.ones((3, 5))

how can i create an boolean mask for this array for all elements of a row up to a specific index stored row-wise in another np.array. So the "end index" array looks like

np.array([0, 2, 3])

The output should look like:

[[True False False False False]
 [True True True False False]
 [True True True True False]]

I tried fancy indexing with a row and column array, but dont know how to specify the column ranges for each row.

Is your input array always an array of 1s? is the first dimension always the size of the mask? — mozway
– mozway, Commented Aug 2, 2024 at 8:36
No the np.ones was just for simplicity, it is a data array. the mask should has the same shape as the data array. i want to multiply the data and mask to zero out values above a certain col index — Netzer
– Netzer, Commented Aug 2, 2024 at 8:51

mozway · Accepted Answer · 2024-08-02 08:36:57Z

1

You don't really need the array of ones, you just need the second dimension (5, which can be inferred from the shape of data), then perform broadcasting with arange and your mask:

N = 5
# or N = data.shape[1]

mask = np.array([0, 2, 3])

np.arange(N) <= mask[:, None]

Output:

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

edited Aug 2, 2024 at 8:36

answered Aug 2, 2024 at 8:26

mozway

267k13 gold badges56 silver badges106 bronze badges

Sign up to request clarification or add additional context in comments.

Collectives™ on Stack Overflow

numpy array mask with range of elements specified by another array row-wise

1 Answer 1

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Comments

Your Answer

Sign up or log in

Post as a guest

Related