如何高效获取numpy.unique返回的有序逆索引？

默认numpy.unique(return_inverse=True)返回的逆索引基于排序后的唯一值，但我们需要的是基于元素首次出现顺序的逆索引，示例对比：

原代码及默认输出：

import numpy as np

arr = np.array([2, 2, 3, 1])
arr_unq, arr_inv = np.unique(arr, return_inverse=True)
print(arr_inv)  # 输出: [1 1 2 0]

期望输出：[0 0 1 2]（对应唯一值顺序：2, 3, 1）

下面针对10万元素、1万唯一值的规模，提供两种高效实现方案：

方法一：纯numpy矢量化实现 #

利用np.unique的return_index参数结合索引映射，全程无Python循环，性能最优：

import numpy as np

arr = np.array([2, 2, 3, 1])

# 获取唯一值、首次出现索引、默认逆索引
_, first_occur_idx, arr_inv_default = np.unique(arr, return_index=True, return_inverse=True)

# 按首次出现的位置排序，得到唯一值的顺序索引
sorted_order = np.argsort(first_occur_idx)

# 创建映射表：将默认逆索引转换为首次出现顺序的索引
inv_map = np.zeros_like(sorted_order)
inv_map[sorted_order] = np.arange(len(sorted_order))

# 生成目标逆索引
arr_inv_ordered = inv_map[arr_inv_default]

print(arr_inv_ordered)  # 输出: [0 0 1 2]

原理说明： #

1. first_occur_idx记录每个唯一值在原数组中第一次出现的位置，示例中为[3, 0, 2]（对应排序后的唯一值[1,2,3]）
2. np.argsort(first_occur_idx)得到按首次出现顺序排列的唯一值索引：[1,2,0]，对应唯一值顺序[2,3,1]
3. inv_map完成索引映射，将默认逆索引的数值转换为目标顺序的索引值

方法二：pandas简洁实现 #

如果项目已依赖pandas，pd.factorize可以直接返回按首次出现顺序的逆索引，代码更简洁：

import pandas as pd
import numpy as np

arr = np.array([2, 2, 3, 1])
arr_inv_ordered, arr_unq_ordered = pd.factorize(arr)

print(arr_inv_ordered)  # 输出: [0 0 1 2]
print(arr_unq_ordered)  # 输出: [2 3 1]

pd.factorize内部为矢量化实现，同样适配大规模数据场景。

内容的提问来源于stack exchange，提问作者Amin.A