如何按首元素小于3的条件对numpy三维数组对应行整行清零

核心逻辑：先提取所有行的首个元素生成布尔判断掩码，再通过 numpy 布尔索引直接对符合条件的行批量置0。
完整可运行代码如下：

import numpy as np

# 初始化输入的3维数组
input_arr = np.array([
    [[0, 2, 3, 4],
     [4, 2, 3, 4],
     [6, 2, 3, 4]],
    [[2, 2, 3, 4],
     [3, 2, 3, 4],
     [2, 2, 3, 4]]
])

# 生成掩码：判断每行第一个元素是否小于3
mask = input_arr[:, :, 0] < 3
# 符合条件的行所有元素置0
input_arr[mask] = 0

print(input_arr)

按你描述的规则执行后输出为：

[[[0 0 0 0]
  [4 2 3 4]
  [6 2 3 4]]
 [[0 0 0 0]
  [3 2 3 4]
  [0 0 0 0]]]

如果你需要得到你给出的示例期望输出（第二个二维数组所有行全置0），只需要将判断条件修改为 input_arr[:, :, 0] <= 3 即可。

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