如何无显式循环按条件高效替换numpy.ndarray数组元素

我有一个形状为(5, 4, 3)的三维numpy.ndarray，需要将数组内部所有长度为3、取值为[3, 3, 3]的子数组替换为[0, 0, 0]。目前找到的实现方案效率不足，处理大规模数组时耗时极长。

测试用数组构造代码如下：

import numpy as np
c = np.array([[[3,3,3],[2,2,2],[3,3,3],[4,4,4]],[[1,1,1],[2,2,2],[7,3,3],[4,4,4]],[[1,1,1],[3,3,3],[3,3,3],[4,4,4]],[[1,1,1],[2,2,2],[3,8,3],[3,3,3]],[[3,3,3],[2,2,2],[3,3,3],[4,4,4]]])

打印数组和形状的输出为：

>>> c
array([[[3, 3, 3],
        [2, 2, 2],
        [3, 3, 3],
        [4, 4, 4]],

       [[1, 1, 1],
        [2, 2, 2],
        [7, 3, 3],
        [4, 4, 4]],

       [[1, 1, 1],
        [3, 3, 3],
        [3, 3, 3],
        [4, 4, 4]],

       [[1, 1, 1],
        [2, 2, 2],
        [3, 8, 3],
        [3, 3, 3]],

       [[3, 3, 3],
        [2, 2, 2],
        [3, 3, 3],
        [4, 4, 4]]])
>>> c.shape
(5, 4, 3)

原有实现代码如下，虽然结果正确，但依赖Python层循环逐行判断，性能极差：

# 将三维数组重构为二维
c_copy = c.reshape(c.shape[0] * c.shape[1], c.shape[2])
# 循环逐行判断是否为目标值，执行替换
c_copy[[np.array_equal(e, [3, 3, 3]) for e in c_copy]] = [0,0,0]
# 重构回原数组形状
c_modified = c_copy.reshape(c.shape)

该实现可以得到正确的替换结果：

>>> c_modified
array([[[0, 0, 0],
        [2, 2, 2],
        [0, 0, 0],
        [4, 4, 4]],

       [[1, 1, 1],
        [2, 2, 2],
        [7, 3, 3],
        [4, 4, 4]],

       [[1, 1, 1],
        [0, 0, 0],
        [0, 0, 0],
        [4, 4, 4]],

       [[1, 1, 1],
        [2, 2, 2],
        [3, 8, 3],
        [0, 0, 0]],

       [[0, 0, 0],
        [2, 2, 2],
        [0, 0, 0],
        [4, 4, 4]]])

但Python层循环的执行效率极低，处理总元素量达900万的大规模数组时耗时无法接受。

完全使用numpy的向量化操作替代Python层循环，所有计算都在C底层执行，性能可以提升2~3个数量级，代码也更简洁，不需要额外做维度重构：

# 沿最后一个轴判断：子数组所有元素都等于3的位置标记为True，得到和前两维形状一致的布尔掩码
mask = (c == 3).all(axis=-1)
# 直接通过掩码批量替换所有匹配的子数组
c[mask] = [0, 0, 0]

如果需要保留原始数组不被修改，提前执行c_copy = c.copy()，对副本执行上述替换操作即可。

性能说明 #

• 原有实现的逐行判断、列表推导逻辑都运行在Python解释器层，循环调度开销极大，完全无法发挥numpy的性能优势
• 向量化实现全程为numpy底层C实现的连续内存操作，对于900万元素规模的数组，通常仅需数毫秒即可完成计算，完全满足大规模数据处理需求。

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