How to index a numpy 2d array with a 2d array mask

Question

Let's say we have a numpy 2D array like the following:

x = array([[0, 7, 1, 6, 2, 3, 4],
           [4, 5, 0, 1, 2, 7, 3]])

and a 2D mask like the following:

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

I'm trying to use the mask in order to get the elements for each row. So the output should look something like this:

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

If I use x[mask] I get array([0, 1, 2]) which is wrong because it flattens out all the selected items.

Any ideas to return it as a 2D array?

These kinds of operations are not well-defined for numpy arrays. I would suggest taking a look at the Awkward Array library — Kevin
– Kevin, Commented Mar 10, 2022 at 13:14
your desired result is not a 2d array! What do you expect to do with the result? — hpaulj
– hpaulj, Commented Mar 10, 2022 at 14:18
numpy.org/doc/stable/user/… elaborates on boolean array indexing. Since, in general, this kind of indexing does not produce the same number of elements in any one dimension, the result is flat. — hpaulj
– hpaulj, Commented Mar 10, 2022 at 16:28
@hpaulj I'm trying to create a generalized version of numpy's intersect1d: numpy.org/doc/stable/reference/generated/numpy.intersect1d.html — Giorgos Perakis
– Giorgos Perakis, Commented Mar 10, 2022 at 18:32

Johan · Accepted Answer · 2022-03-10 11:56:54Z

2

How about

[xi[mi] for xi,mi in zip(x,mask)]

answered Mar 10, 2022 at 11:56

Johan

5194 silver badges10 bronze badges

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

1 Comment

I thought about that too, but I want to avoid doing a for loop. I want to do it in a more vectorized way. Thanks!