numpy argsort排序相等元素时如何保留原有相对顺序

云上老周

2026-8-29

问题说明

现有如下numpy数组：

array([0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00,
    0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00,
    0.00000000e+00, 0.00000000e+00, 0.00000000e+00, 0.00000000e+00,
    0.00000000e+00, 0.00000000e+00, 1.78772666e+02, 6.35416628e+02,
    1.16901104e+03, 1.79252306e+03, 2.52110498e+03, 3.37246243e+03,
    4.36728470e+03, 5.52974748e+03, 6.88810037e+03, 8.47535341e+03,
    1.03300792e+04, 1.24973505e+04, 1.50298353e+04, 1.79890767e+04,
    2.14469887e+04, 2.54876040e+04, 3.02091156e+04, 3.57262634e+04,
    4.21731224e+04, 4.97063611e+04, 5.85090485e+04, 6.87951035e+04,
    8.08144926e+04, 9.48593053e+04, 1.11270852e+05, 1.30447958e+05,
    1.52856655e+05, 1.79041511e+05, 2.09638855e+05, 2.45392250e+05,
    2.87170558e+05, 3.35989054e+05, 3.93034103e+05, 4.59691985e+05,
    5.37582589e+05, 6.28598773e+05]))

直接调用默认argsort()返回的索引中，前14个值为0的元素索引顺序混乱，结果如下：

array([ 0, 13, 12, 10,  9,  8,  7, 11,  5,  4,  3,  2,  1,  6, 14, 15, 16,
    17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33,
    34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49],
   dtype=int64)

需求为：数组元素值相等时，排序后的索引按原数组位置递增排列，保留相等元素的原始相对顺序，实现稳定排序效果。

解决方案

默认argsort()使用的是快速排序（quicksort）算法，该算法不属于稳定排序，相等元素的相对顺序不做保证，因此会出现顺序混乱的问题。
只需要在调用argsort时指定排序算法为稳定排序即可，这是性能最优、代码最简洁的方案，不需要额外构造辅助数组做二次处理：

推荐写法（numpy 1.15及以上版本支持）：指定kind='stable'，numpy会自动选择当前版本下性能最优的稳定排序实现，时间复杂度为O(nlogn)，和默认快速排序性能差距极小。
示例代码：
```
# arr为你的目标数组
sorted_indices = arr.argsort(kind='stable')
```
兼容旧版本写法：指定kind='mergesort'，归并排序是经典稳定排序算法，所有numpy版本都支持，排序效果和stable完全一致。
示例代码：
```
sorted_indices = arr.argsort(kind='mergesort')
```

针对你给出的测试数组，上述代码返回的结果中，前14个0值元素的索引会严格按照0,1,2,...,13的顺序排列，非相等元素的排序逻辑和默认argsort完全一致，符合需求。

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