NumPy按条件替换ndarray整行数组的高效实现方案

本问题聚焦于满足条件时对ndarray中的整行（整个子数组）进行替换，而非ndarray内单个元素值的替换——单元素替换场景可参考对应单元素条件替换的相关解答。

问题描述 #

现有输入数组 inarray，示例如下：

import numpy as np
inarray = \
np.array([[[1,1,1,1],
           [2,2,2,2]],
          [[3,3,3,3],
           [4,4,4,4]]])

需求为将所有内容为 [1,1,1,1] 的行整体替换为 [9,9,9,9]，需要比现有实现更简洁高效的方案。

已验证的基础实现 #

当前自行编写的可生效代码如下：

inarray[np.equal(inarray,np.array([1,1,1,1])).all(axis=2)]=np.array([9,9,9,9])

该行代码执行后，inarray 修改结果为：

array([[[9, 9, 9, 9],
        [2, 2, 2, 2]],

       [[3, 3, 3, 3],
        [4, 4, 4, 4]]])

完全符合预期效果。

优化实现方案 #

上述实现逻辑完全正确，可进一步简化写法，在不损失运行效率的前提下提升代码可读性：

inarray[(inarray == [1,1,1,1]).all(axis=2)] = [9,9,9,9]

• 优化点说明：直接使用==运算符完成数组逐元素比较，和np.equal效果完全一致，无需额外显式转换目标数组为ndarray类型，numpy会自动通过广播机制完成维度匹配，代码更简洁。
• 适配性说明：该写法无论匹配到多少行符合条件的子数组，都可以正确完成整行赋值，不会出现维度不匹配的报错。

常见错误方案说明 #

有公开解答建议使用如下写法实现替换：

inarray[inarray == [1,1,1,1]] = [9,9,9,9]

该写法存在本质缺陷：inarray == [1,1,1,1] 生成的是逐元素匹配的布尔掩码，掩码为True的位置是所有值为1的单个元素，而非整行匹配的位置。当数组中存在多个 [1,1,1,1] 行时，掩码为True的元素数量和待赋值的数组长度不匹配，会直接运行报错，测试代码如下：

inarray = \
np.array([[[1,1,1,1],
           [2,2,2,2]],
          [[1,1,1,1],
           [4,4,4,4]]])

inarray[inarray == [1,1,1,1]] = [9,9,9,9]

运行后抛出如下错误：

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Input In [81], in <cell line: 1>()
----> 1 inarray[inarray == [1,1,1,1]] = [9,9,9,9]

ValueError: NumPy boolean array indexing assignment cannot assign 4 input values to the 8 output values where the mask is true

该测试场景下的预期运行结果为 inarray 修改为：

array([[[9, 9, 9, 9],
        [2, 2, 2, 2]],

       [[9, 9, 9, 9],
        [4, 4, 4, 4]]])

使用上述优化后的实现即可正确得到该结果。

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