如何在for循环中高效为NumPy数组元素分配递增索引值

如果追求极致性能，NumPy的向量化操作本身比手动for循环高效得多，但既然你明确要基于for循环实现，下面是两种最高效的思路：

方法一：字典映射+索引遍历（最优） #

先构建值到目标索引的映射字典（字典查找是O(1)时间复杂度），然后通过遍历数组索引的方式赋值——直接遍历NumPy数组元素会产生额外的对象转换开销，用索引访问能大幅减少这部分损耗。

示例代码：

import numpy as np

# 示例原始数组
original_arr = np.array(['A', 'B', 'C', 'D', 'A', 'C', 'B'])
# 预定义值到递增数字的映射
value_to_idx = {'A': 1, 'B': 2, 'C': 3, 'D': 4}

# 初始化结果数组，和原数组形状一致，指定int类型
result_arr = np.empty_like(original_arr, dtype=np.int32)

# 高效遍历索引赋值
for idx in range(original_arr.size):
    result_arr[idx] = value_to_idx[original_arr[idx]]

print(result_arr)
# 输出：[1 2 3 4 1 3 2]

方法二：利用唯一值匹配（适合值不固定的场景） #

如果你的预定义值不是硬编码的，而是从数组中提取的唯一值，可以先获取排序后的唯一值，再通过循环匹配索引：

import numpy as np

original_arr = np.array(['A', 'B', 'C', 'D', 'A', 'C'])
# 获取排序后的唯一值，对应1、2、3、4的顺序
unique_vals = np.unique(original_arr)
# 构建映射字典（自动对应递增数字）
value_to_idx = {val: i+1 for i, val in enumerate(unique_vals)}

result_arr = np.empty_like(original_arr, dtype=np.int32)
for idx in range(original_arr.size):
    result_arr[idx] = value_to_idx[original_arr[idx]]

额外提示 #

如果不限制必须用for循环，更推荐直接用向量化操作，比如：

result_arr = np.vectorize(value_to_idx.get)(original_arr)

或者更高效的列表推导式转数组：

result_arr = np.array([value_to_idx[val] for val in original_arr], dtype=np.int32)

但手动for循环的话，字典映射+索引遍历是性能最优的方案。

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