NumPy 如何无循环按子数组非零条件将三维数组降为二维数组

输入与转换规则 #

待处理的原始三维NumPy数组结构如下：

import numpy as np
origin_arr = np.array([
    [[0, 1],
     [0, 0]],
    [[0, 0],
     [6, 7]]
])

需要输出的目标二维数组结构为：

array([[1, 0],
       [0, 1]])

转换判定规则：原数组最内层的每个2元素子数组，只要包含任意非零值，对应输出位置标记为1；子数组所有元素均为0时，对应位置标记为0，实现过程不使用手写Python循环。

高效实现方案 #

直接利用NumPy的向量化聚合接口即可完成全流程操作，无需手写遍历：

• 沿数组最后一个轴向（即最内层子数组所在维度）执行存在性判断：调用np.any()逐子数组检查是否存在非零元素，逐位置返回布尔值（存在非零值返回True，全零返回False）
• 将得到的布尔矩阵直接转换为整数类型，即可得到0/1标记的最终结果
  核心实现代码：

result = np.any(origin_arr, axis=-1).astype(int)

该实现完全运行在NumPy底层的C级计算路径上，无Python层的循环开销，哪怕处理高维度、大规模的数组数据，也能保持极高的运算效率。

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