CODEWITHSUNDEEP
pythonnumpy
incognito
incognito

Reputation: 199

Mask an 2Darray row wise by another array

Consider the following 2d array:

>>> A = np.arange(2*3).reshape(2,3)
array([[0, 1, 2],
       [3, 4, 5]])

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

I would like to get the following mask from A as row wise condition from b as an upper index limit:

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

Can numpy do this in a vectorized manner?

Upvotes: 1

Views: 57

Answers (2)

PaulS
PaulS

Reputation: 25323

Another possible solution, based on the idea that the wanted mask corresponds to a boolean lower triangular matrix:

mask = np.tril(np.ones(A.shape, dtype=bool))

Output:

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

Upvotes: 0

mozway
mozway

Reputation: 260640

You can use array broadcasting:

mask = np.arange(A.shape[1]) < b[:,None]

output:

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

Upvotes: 3

Related Questions