如何高效基于另一数组替换Numpy数组中的对应元素？

可以通过构建映射字典的方式实现，完全避免循环查找索引，效率远高于np.where循环的方案，具体步骤如下：

1. 基于第二个数组创建元素到其下一个元素的映射字典，利用数组切片和zip快速生成键值对
2. 将第一个数组扁平化后通过字典完成批量映射，最后恢复原形状

代码示例 #

import numpy as np

# 定义输入数组
arr1 = np.array([['CLU20'], ['CLZ20'], ['CLH21']])
arr2 = np.array(['CLU20', 'CLZ20', 'CLH21', 'CLM21', 'CLU21', 'CLZ21', 'CLH22', 'CLM22', 'CLU22'])

# 构建元素到下一个元素的映射字典
next_map = dict(zip(arr2[:-1], arr2[1:]))

# 完成替换：扁平化arr1 -> 字典映射 -> 恢复原形状
result = np.array([next_map[x] for x in arr1.flatten()]).reshape(arr1.shape)

print(result)

输出结果 #

[['CLZ20']
 ['CLH21']
 ['CLM21']]

优势说明 #

• 字典的查找时间复杂度为O(1)，相比循环调用np.where（每次查找O(n)），在数组规模较大时性能提升非常明显
• 代码简洁直观，逻辑清晰，避免了繁琐的索引计算

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