如何优化Python中3D numpy数组阈值过滤的迭代实现代码？

代码优化方案 #

你写的嵌套循环代码有两个核心问题：

1. Python层级的循环完全浪费了numpy的向量化运算优势，所有判断、赋值操作都在解释器层面执行，速度极慢
2. 代码里b = a属于浅拷贝，修改b的同时会直接修改原始输入数组a，如果后续还要用到原数组会出现非预期的结果

下面是几种可直接使用的优化方案，全部基于numpy的C级向量化运算，速度比你原有代码快100~1000倍不等，数组规模越大性能优势越明显，且自动兼容任意维度的输入数组，不需要硬编码循环层数。

方案1：np.where实现（可读性最高） #

逻辑清晰，返回新数组不会修改原始输入：

import numpy as np

def array_max(a):
    return np.where(a > 1.4, a, 0)

方案2：布尔掩码乘法（写法最简洁） #

利用numpy广播机制，布尔值True会被转为1、False转为0，相乘直接得到结果：

def array_max(a):
    return a * (a > 1.4)

方案3：原地修改（内存占用最低） #

如果不需要保留原始输入数组的内容，可直接在原数组上修改，省去内存拷贝开销：

def array_max(a):
    a[a <= 1.4] = 0
    return a

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