如何用列表推导式高效替换NumPy数组中的指定元素

直接利用NumPy的布尔索引和向量化赋值就能高效完成需求，完全不需要Python循环，代码如下：

import numpy as np

nparray1 = np.array([1, 1, 0, 0, 1])
nparray2 = np.array([8,7,4])

# 生成标记数组：nparray1中值为1的位置对应True
mask = nparray1 == 1
# 批量替换对应位置的元素
nparray1[mask] = nparray2

print(nparray1)  # 输出: [8 7 0 0 4]

方法说明 #

• 第一步通过nparray1 == 1生成布尔数组mask，精准定位所有需要替换的位置；
• 利用NumPy的布尔索引完成批量赋值，这是底层优化的向量化操作，比Python循环快得多，数组规模越大，效率优势越明显。

原循环方法的不足 #

原代码的循环逻辑每次切片nparray2 = nparray2[1:]会生成新数组对象，且Python级别的循环在处理大数据量时性能远低于NumPy的向量化操作。

关于列表推导式的补充 #

如果想用列表推导式实现，需要额外维护指针跟踪nparray2的当前位置，示例如下：

import numpy as np

nparray1 = np.array([1, 1, 0, 0, 1])
nparray2 = np.array([8,7,4])

ptr = 0
result = []
for x in nparray1:
    if x == 1:
        result.append(nparray2[ptr])
        ptr += 1
    else:
        result.append(x)
result = np.array(result)

但这种写法本质还是循环，效率和原代码相近，远不如NumPy的向量化操作高效，不推荐使用。

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