如何用Numpy获取元素排序后的位置索引而非排序数组？

问题背景 #

给定3D数组 arr = [[[20, 5, 10], ...], ...]，调用 arr.argsort(axis=-1) 得到的是沿最后一个轴的排序索引（即每个位置的值对应原数组中该轴上的元素下标，示例结果为[[[1, 2, 0], ...], ...]），但实际需求是获取原数组每个元素在排序后数组中对应的位置索引（示例期望结果为[[[2, 0, 1], ...], ...]）。

解决方案 #

直接对argsort的结果再次调用argsort(axis=-1)即可，这是利用了argsort的逆操作特性，完全基于Numpy向量化实现，无需列表推导或循环，适合大规模数组处理。

示例代码 #

import numpy as np

# 构造示例3D数组
arr = np.array([[[20, 5, 10], [15, 30, 25]]])

# 获取沿最后轴的排序索引
sort_indices = arr.argsort(axis=-1)
# 对排序索引再次执行argsort，得到原元素在排序后的位置索引
reverse_indices = sort_indices.argsort(axis=-1)

print("原数组：")
print(arr)
print("\n排序索引（arr.argsort()结果）：")
print(sort_indices)
print("\n原元素在排序后的位置索引：")
print(reverse_indices)

原理说明 #

• arr.argsort(axis=-1) 返回的是将原数组沿最后轴排序时，排序后数组每个位置对应的原数组下标。比如原数组[20,5,10]排序后为[5,10,20]，argsort结果[1,2,0]表示排序后第0位取原数组第1位元素，第1位取原数组第2位元素，第2位取原数组第0位元素。
• 对该排序索引数组再次执行argsort(axis=-1)，相当于反向映射：排序索引数组中的值代表原数组下标，它们在排序索引数组中的位置就是该原下标元素在排序后的位置。例如[1,2,0]的argsort结果是[2,0,1]，对应原数组第0位元素（20）在排序后是第2位，第1位元素（5）在排序后是第0位，第2位元素（10）在排序后是第1位，完全符合需求。

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