如何基于NumPy数组生成原数组元素排序后的索引数组？

无重复元素/保留重复元素原顺序的高效方案 #

使用两次argsort即可快速得到需求结果：

import numpy as np

arr = np.array([10, 5, 13, 2])
sorted_indices = arr.argsort()
result = sorted_indices.argsort()

print(result)  # 输出：array([2, 1, 3, 0])

原理 #

• 第一次arr.argsort()返回的是排序后元素对应的原数组索引，示例中为[3, 1, 0, 2]，代表排序后的数组元素依次来自原数组的第3、1、0、2位。
• 对该结果再次调用argsort()，本质是查找原数组每个索引在sorted_indices中的位置——这个位置就是原数组对应元素在排序后数组中的索引。

重复元素统一映射到首次出现索引的方案 #

如果数组包含重复元素，且希望相同元素都映射到排序后首次出现的索引，可结合np.sort和np.searchsorted：

import numpy as np

arr = np.array([5, 3, 5, 2])
sorted_arr = np.sort(arr)
result = np.searchsorted(sorted_arr, arr)

print(result)  # 输出：array([2, 1, 2, 0])

np.searchsorted会在排序后的数组中查找原数组元素的插入位置，默认返回第一个大于等于目标元素的位置，恰好对应重复元素首次出现的索引。

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