如何用Numpy的where/布尔索引实现符合指定条件的三维数组切片？

问题说明 #

现有如下Numpy三维数组a：

import numpy as np

a = np.array([[[3, 2, 2], [3, 4, 2]],
              [[1, 2, 2], [3, 4, 2]],
              [[1, 2, 2], [3, 4, 2]]
])
print(a)

输出结果：

[[[3 2 2]
  [3 4 2]]

 [[1 2 2]
  [3 4 2]]

 [[1 2 2]
  [3 4 2]]]

需求是筛选出每个二维子数组的[0,0]和[1,0]元素均大于2的子数组。目前通过循环可以实现：

b = []
for i in range(a.shape[0]):
    if (a[i, 0, 0] > 2 and a[i, 1, 0] > 2):
        b.append(a[i])
print(np.array(b))

输出结果：

[[[3 2 2]
  [3 4 2]]]

但尝试以下两种向量化索引方法时，得到了不符合需求的结果：
方法1：

a[np.where(a[:,:,0] > 2)]

方法2：

a[a[:,:,0]> 2]

两种方法的输出均为：

array([[3, 2, 2],
       [3, 4, 2],
       [3, 4, 2],
       [3, 4, 2]])

需要找到正确的向量化索引方法实现需求。

正确的向量化实现 #

可以通过构造一维布尔掩码来筛选三维数组的第一个维度（即每个二维子数组），具体有两种实现方式：

方式一：直接判断两个位置的条件 #

先分别判断每个二维子数组的[0,0]和[1,0]是否大于2，再取逻辑与得到掩码：

# 构造掩码：每个二维子数组的[0,0]和[1,0]都大于2
mask = (a[:, 0, 0] > 2) & (a[:, 1, 0] > 2)
# 用掩码索引原数组
result = a[mask]
print(result)

输出结果：

[[[3 2 2]
  [3 4 2]]]

方式二：利用all()方法批量判断 #

先提取每个二维子数组的第一列（即所有[x,0]位置的元素），再通过all(axis=1)判断该子数组的所有行的第一列元素是否都大于2：

# 提取每个二维子数组的第一列，判断是否全大于2，得到一维掩码
mask = (a[:, :, 0] > 2).all(axis=1)
# 索引得到结果
result = a[mask]
print(result)

输出结果与方式一一致。

错误方法的原因 #

之前的错误方法中，a[:,:,0] > 2生成的是一个(3,2)的二维布尔数组，直接用它索引原三维数组时，Numpy会将其视为提取所有满足条件的二维行元素，而不是筛选整个二维子数组，因此得到的是零散的行，而非完整的目标子数组。

内容的提问来源于stack exchange，提问作者yangtzech