如何正确掩码NumPy数组 提取满足条件值并保持原尺寸零填充

针对提取pe数组中值不等于255的元素、同形状输出、不满足条件位置填0的需求，最优实现完全基于NumPy原生向量化操作，无Python层循环开销，性能达到最优。

核心实现代码 #

import numpy as np

pd = np.random.randint(254,256, size=(7,7))
pe = np.random.randint(0,7, size=(7,7))

# 掩码核心逻辑：pe!=255时保留pe原值，否则填充0
result = np.where(pe != 255, pe, 0)

实现说明 #

• 性能优势：np.where是NumPy内置C级实现的函数，直接在内存块层面完成条件判断和赋值，没有Python解释器的循环调度开销，是这类同尺寸掩码赋值场景性能最高的写法。
• 尺寸匹配：输出结果自动和输入pe数组形状保持一致，无需额外做维度调整，直接满足(7,7)的形状要求。
• 结果验证：执行后result中所有pe值为255的位置都会被填充为0，其余位置和pe对应位置取值完全一致。

等价低内存写法（可选） #

如果追求极致内存占用，且不需要保留原pe数组的原始值，可以用原地赋值的写法，性能和np.where几乎一致：

result = pe.copy()
result[pe == 255] = 0

注意：不要使用Python层面的for循环遍历、列表推导等方式实现掩码操作，这类写法在数组尺寸较大时，性能会比原生向量化实现低2个数量级以上。

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