如何使用Numpy将数组重复元素替换为其首次出现的索引

核心实现思路 #

利用np.unique的内置参数直接获取首次出现位置和反向映射关系，全程为向量化操作，效率远高于Python原生循环方案。

代码示例 #

import numpy as np

# 输入数组
x = np.array([1, 2, 3, 3, 2, 4])
# 自定义索引数组，如果是默认0到n-1的索引可省略直接用first_idx
I = np.array([0, 1, 2, 3, 4, 5])

# 参数说明：
# return_index：返回每个唯一元素在原数组中第一次出现的下标
# return_inverse：返回原数组每个元素对应到唯一值数组的下标映射
_, first_idx, inverse_map = np.unique(x, return_index=True, return_inverse=True)

# 生成结果：用映射下标取对应首次出现的索引
y = I[first_idx][inverse_map]

print(y)
# 输出：[0 1 2 2 1 5]

方案优势 #

• 无需手动遍历数组，完全依赖Numpy底层C实现的运算，处理十万级以上数组时速度是Python字典遍历方案的10~100倍
• 代码简洁易读，仅需2行核心逻辑即可完成需求
• 支持自定义索引数组，适配非连续索引的特殊场景

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