Filter numpy array base 1-0 掩码

Question

Say I have a numpy ndarray of shape (172,40,20) and a 1-0 mask of shape (172, 40). I would like to do something similar to bitwise_and: to keep those values where the mask value is 1 while other values set to 0 where the mask value is 0.

Is there a way that I can do without looping? for example

a3D = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
print(a3D.shape)
mask = np.array([[1, 0], [0, 1]])
print(mask.shape)
# (2, 2, 2)
# (2, 2)

# my func
result = bitwise_and(a3D, mask) 

# result = np.array([[[1, 2], [0, 0]], [[0, 0], [7, 8]]])

Answer 1

Simply index the array using the mask as boolean indices :

result = a3D.copy()
result[mask == 0, :] = 0
print(result)

outputs

[[[1 2]
  [0 0]]

 [[0 0]
  [7 8]]]

Answer 2

You can use np.where() :

np.where(mask[:, :, np.newaxis], a3D, 0)

[[[1 2]
  [0 0]]

 [[0 0]
  [7 8]]]