[Python] Index array using boolean mask with a "broadcasted" dimension?

I have two arrays a1 and a2 of shape M1xM2x...xMPx3 and N1xN2x...xNQx3, and a mask m (boolean array) of shape M1xM2x...xMPxN1xN2x...xNQ (you can assume that a1 and a2 are at least 2D).

import numpy as np

np.random.seed(0)

a1 = np.random.rand(4, 5, 3)
a2 = np.random.rand(6, 3)

m = np.random.rand(4, 5, 6) >= 0.7


I would like to two arrays b1 and b2 of shape Mx3 that contains the values of a1 and a2 where the mask m is true.

My current way is to repeat a1 and a2 to obtain arrays of size M1xM2x...xMPxN1xN2x...xNQx3 and then index them using m:

a1_shape, a2_shape = a1.shape[:-1], a2.shape[:-1]
t_shape = a1_shape + a2_shape + (3,)

a1 = a1[..., None, :].repeat(np.prod(a2_shape), axis=-2).reshape(t_shape)
a2 = a2[None, ..., :].repeat(np.prod(a1_shape), axis=0).reshape(t_shape)

b1, b2 = a1[m, :], a2[m, :]


This requires creating two huge array of size M1xM2x...xMPxN1xN2x...xNQ which can be problematic.

Is there a way to obtain the same b1 and b2 without having to create such huge intermediate array?

Continue reading...