使用2d布尔从3d numpy数组中提取1d数组

Question

假设我有一个3D numpy数组：

i, j, k = 10, 3, 4

arr = np.arange(120).reshape(i, j, k)

和一个二维布尔数组：

mask = np.random.random((j, k)) > 0.5

n = mask.sum()

我希望能够从arr沿其1st True取与mask的True值相对应的1d数组。 结果应具有形状(i, n) 。 怎么办呢？

我拉出一些旧代码，由于某种原因我在做arr[mask]但这给出了(n, k)的形状（我不确定为什么）和一个警告：

VisibleDeprecationWarning: boolean index did not match indexed array along dimension 0; dimension is 10949 but corresponding boolean dimension is 11

Answer 1

只需沿最后两个轴遮罩-

arr[:,mask]