使用 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 沿其第一维提取一维数组，该维对应于 mask 的 True 值。结果应具有 (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]