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?

Answer 1

How about

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