基于预计算排序数组的查询元素排序实现（Numpy）

需求描述 #

给定如下无序数组：

import numpy as np
main_list = np.array([100,200,400,1000,800,900,700,600,500,300])

以及查询索引列表 q = np.array([2,5,7,9])（对应元素为 np.array([400, 900, 600, 300])），需要直接获取main_list中对应元素的排序结果及原位置，期望输出：

np.array([300, 400, 600, 900])
np.array([9,2,7,5])

核心要求：仅使用预计算的排序数组，避免每次查询时执行排序操作。

Numpy解决方案（感谢@cary-swoveland） #

离线预计算步骤 #

预先计算两个排序索引数组，后续查询直接复用：

argsort = np.argsort(main_list)
argsort_argsort = np.argsort(argsort)

查询阶段 #

基于预计算的数组，快速得到排序后的原索引和对应元素：

q = np.array([2,5,7,9])

new_array = np.full(main_list.shape[0], -1)
new_array[argsort_argsort[q]] = q

sorted_q = new_array[new_array != -1]
sorted_values = main_list[sorted_q]

print(sorted_q)
print(sorted_values)

执行结果 #

运行上述代码后，输出如下：

[9 2 7 5]
[300 400 600 900]

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