如何使用np.unique()按指定顺序排序分类变量的唯一值？

我尝试用以下代码获取分类变量的唯一值并按指定顺序排序，但没成功：

import numpy as np

unique_values, unique_value_counts = np.unique(['Small', 'Medium', 'Large', 'Medium', 'Small', 'Large', 'Small', 'Medium'], return_counts = True)

print(unique_values)

代码输出：

['Large', 'Medium', 'Small']

但我期望得到按Small、Medium、Large顺序排列的结果：

['Small', 'Medium', 'Large']

请问是否可以通过np.unique()实现该排序需求？

解答 #

np.unique()默认按字典序排序返回结果，所以会得到['Large', 'Medium', 'Small']。该函数本身没有直接指定自定义排序规则的参数，但可以结合其他步骤实现需求：

方法1：利用有序分类类型（推荐） #

借助pandas.Categorical定义自定义顺序，再传入np.unique()处理：

import numpy as np
import pandas as pd

# 指定目标顺序
custom_order = ['Small', 'Medium', 'Large']
# 将原始数组转为有序分类
categorical_arr = pd.Categorical(
    ['Small', 'Medium', 'Large', 'Medium', 'Small', 'Large', 'Small', 'Medium'],
    categories=custom_order,
    ordered=True
)
# 获取唯一值和计数
unique_vals, counts = np.unique(categorical_arr, return_counts=True)
print(unique_vals)

输出：

['Small' 'Medium' 'Large']

方法2：手动映射排序 #

如果不想依赖pandas，可以先获取默认结果，再按自定义顺序重新提取：

import numpy as np

custom_order = ['Small', 'Medium', 'Large']
original_arr = ['Small', 'Medium', 'Large', 'Medium', 'Small', 'Large', 'Small', 'Medium']

# 获取默认去重后的结果和计数
unique_vals, counts = np.unique(original_arr, return_counts=True)
# 构建值与计数的映射
val_count_map = dict(zip(unique_vals, counts))
# 按自定义顺序生成结果
sorted_vals = np.array(custom_order)
sorted_counts = np.array([val_count_map[val] for val in custom_order])

print(sorted_vals)
print(sorted_counts)

输出：

['Small' 'Medium' 'Large']
[3 3 2]

补充说明 #

np.unique()的核心逻辑是去重后执行默认字典序排序，没有内置自定义排序的参数。上述两种方法都能基于np.unique()的能力实现自定义顺序需求，其中方法1更贴合分类变量的处理场景，代码更简洁。

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