2d boolean selection in 3d matrix

Question

In a related question I learned that if I have an array of shape MxMxN, and I want to select based on a boolean matrix of shape MxM, I can simply do

data[select, ...]

and be done with it. Unfortunately, now I have my data in a different order:

import numpy as np
data = np.arange(36).reshape((3, 4, 3))
select = np.random.choice([0, 1], size=9).reshape((3, 3)).astype(bool)

For each element in data indexed i0, i1, i2, it should be selected, if select[i0, i2] == True.

How can I proceed with my selection without having to do something inefficient like

data.flatten()[np.repeat(select[:, None, :], 4, axis=1).flatten()]

Answer 1

One way would be to simply use np.broadcast_to to broadcast without actual replication and use that broadcasted mask directly for masking required elements -

mask = np.broadcast_to(select[:,None,:], data.shape)
out = data[mask]

Another way and probably faster one would be to get the indices and then index with those. The elements thus obtained would be ordered by axis=1. The implementation would look something like this -

idx = np.argwhere(select)
out = data[idx[:,0], :, idx[:,1]]