Numpy替换三维数组中[1,1,1]子数组为[255,255,255]的高效方法

问题背景 #

待处理的是形状为(高度, 宽度, 3)的三维NumPy数组，示例数据如下：

import numpy as np
blackwhite = np.array([[[0, 0, 0],
    [0, 0, 0]],

   [[0, 0, 0],
    [0, 0, 0]],

   [[1, 0, 0],
    [1, 1, 1]]])

需求为将数组中所有值完全等于[1, 1, 1]的3元素子数组，整体替换为[255, 255, 255]，需要适配8192×8192级别的大规模数组，满足性能要求。
此前尝试的两类方案均存在缺陷：

• 多层for循环逐索引判断：代码存在笔误（条件判断处误将比较运算符==写为赋值运算符=，导致最初运行不生效），且Python层面循环处理大数组效率极低，无法满足8K分辨率数据的处理需求。
• 直接调用np.where(blackwhite == [1, 1, 1])：该写法会触发广播机制逐元素匹配值为1的位置，返回的是单个元素的匹配索引，无法定位整个子数组完全匹配的坐标，不能实现整组替换。

最优实现方案 #

全程使用NumPy向量化运算，无Python层面循环，处理8192×8192尺寸数组耗时仅毫秒级，代码如下：

# 生成形状为(高度, 宽度)的布尔掩码，标记所有最后一维完全匹配[1,1,1]的位置
match_mask = (blackwhite == [1, 1, 1]).all(axis=-1)
# 通过布尔索引批量赋值
blackwhite[match_mask] = [255, 255, 255]

原理说明 #

1. blackwhite == [1, 1, 1]会利用NumPy广播规则，对每个3元素子数组逐元素和目标值对比，返回和原数组同形状的布尔数组，每个位置标记对应元素是否等于1。
2. .all(axis=-1)沿最后一个轴（即3元素子数组所在的通道轴）做逻辑与运算，只有子数组内3个元素的对比结果全为True时，对应掩码位置才为True，正好筛选出所有完全匹配[1,1,1]的子数组坐标。
3. 布尔索引批量赋值是NumPy原生支持的向量化操作，所有运算在C层面执行，性能远高于Python循环。

效果验证 #

对示例数组执行上述代码后，输出结果如下，[1,0,0]这类部分元素为1的子数组不会被误修改：

array([[[  0,   0,   0],
        [  0,   0,   0]],

       [[  0,   0,   0],
        [  0,   0,   0]],

       [[  1,   0,   0],
        [255, 255, 255]]])

等价写法 #

如果习惯使用np.where，可以用如下写法，性能和上述方案完全一致：

match_mask = np.all(blackwhite == [1, 1, 1], axis=-1)
blackwhite[np.where(match_mask)] = [255, 255, 255]

该方案内存占用低、运行效率高，是NumPy处理这类整组匹配替换场景的标准写法。

内容的提问来源于stack exchange，提问作者Ahmad Uzzam Masood