在NumPy中如何高效获取向量元素排序后的目标索引映射?
在NumPy中高效获取排序的置换映射关系
如果你的需求是拿到原向量里每个元素在排序后数组对应的索引位置,而且不想额外执行实际的排序操作,那用NumPy的argsort两次就能搞定,这也是最省心高效的方式。
先搞懂核心逻辑
先看第一次调用np.argsort(my_vector):它返回的是一个索引数组,告诉你排序后的数组每个元素是从原数组的哪个位置取来的。拿你的示例数组举例:
import numpy as np my_vector = np.array([23, 20, 33, 38, 27]) sort_indices = np.argsort(my_vector) print(sort_indices) # 输出: [1 0 4 2 3]
这个结果的意思是:排序后的第0个元素(20)来自原数组索引1,第1个元素(23)来自原数组索引0,第2个元素(27)来自原数组索引4,以此类推。
但我们要的是反过来:原数组每个元素在排序后数组里的位置。这时候只要对sort_indices再调用一次np.argsort()就行,相当于找原索引在sort_indices里的位置——这个位置就是它在排序后数组的索引:
reverse_map = np.argsort(sort_indices) print(reverse_map) # 输出: [1 0 3 4 2]
咱们逐个验证:
- 原数组索引0的23,在排序后数组里是第1个元素(索引1)→
reverse_map[0] = 1,对得上 - 原数组索引1的20,在排序后数组里是第0个元素(索引0)→
reverse_map[1] = 0,没问题 - 原数组索引2的33,在排序后数组里是第3个元素(索引3)→
reverse_map[2] = 3,正确 - 原数组索引3的38,在排序后数组里是第4个元素(索引4)→
reverse_map[3] = 4,没错 - 原数组索引4的27,在排序后数组里是第2个元素(索引2)→
reverse_map[4] = 2,完全匹配
一步搞定的写法
嫌两步麻烦?直接把两次调用合并成一行就行:
reverse_map = np.argsort(np.argsort(my_vector))
要是需要降序的映射呢?
如果想基于降序排序来获取映射,有两种简单方式:
一种是对原数组取负后再用argsort,然后重复上面的步骤:
# 降序的反向映射 desc_sort_indices = np.argsort(-my_vector) desc_reverse_map = np.argsort(desc_sort_indices) print(desc_reverse_map) # 示例数组输出: [3 4 1 0 2]
验证下:降序排序后的数组是[38,33,27,23,20],原索引0的23在降序数组里是索引3,原索引1的20是索引4,完全对应得上。
为什么这个方法高效?
np.argsort()是NumPy底层优化过的算法(默认用快速排序,需要稳定排序的话可以选归并排序),时间复杂度是O(n log n),两次调用下来总复杂度还是O(n log n),这是获取这种映射关系的最优复杂度——全程不需要实际排序元素,只靠索引计算就能完成,效率拉满。
内容的提问来源于stack exchange,提问作者hanuta98
相关产品推荐
相关产品推荐

