Mask an 2Darray row wise by another array-CodePudding

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?

CodePudding user response：

You can use array broadcasting:

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

output:

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

CodePudding user response：

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