基于初始argsort结果对Numpy数组子集高效排序的索引方法

核心思路 #

既然你已经有了全局数组x的argsort结果i，我们可以通过预构建逆索引数组的方式，用纯索引操作完成任意子集n的排序，完全避免重复排序或集合检查。

逆索引数组inv_i的作用是记录每个原索引在全局排序序列i中的位置——简单说，inv_i[idx]就是idx在i里的下标，这样我们就能快速知道任意元素在全局排序中的相对位置。

代码实现 #

首先预计算逆索引（只需要执行一次）：

import numpy as np

x = np.array([14, 15,  9,  6, 19, 18,  4, 11, 10,  0])
i = np.argsort(x)

# 构建逆索引：inv_i[idx] 表示原索引idx在全局排序数组i中的位置
inv_i = np.zeros_like(i)
inv_i[i] = np.arange(len(i))

然后定义通用的排序函数，处理任意子集n：

def index_sort(subset, global_pos):
    # 根据全局排序位置对当前子集排序
    sorted_order = np.argsort(global_pos)
    return subset[sorted_order]

# 示例测试
n = np.array([2, 5, 7, 8])
x_subset = x[n]
# 获取子集元素在全局排序中的位置
subset_pos = inv_i[n]
# 得到排序后的子集
sorted_x_subset = index_sort(x_subset, subset_pos)
print(sorted_x_subset)  # 输出: array([ 9, 10, 11, 18])

# 对无关数组y执行相同操作
y = np.array([0,1,2,3,4,5,6,7,8,9])
y_subset = y[n]
sorted_y_subset = index_sort(y_subset, subset_pos)
print(sorted_y_subset)  # 输出: array([2, 8, 7, 5])

为什么之前的尝试失败 #

• x[n][i]：i的长度是原数组x的长度，而x[n]已经是子集，索引范围不匹配，直接报错
• x[i[n]]/x[i][n]：这两种方式都是取全局排序数组中n位置的元素，拿到的是全局排序后的第n个元素，不是子集n内部的排序结果，所以数据完全不对

性能优势 #

• 逆索引数组inv_i仅需计算一次，时间复杂度O(N)（N为x的长度）
• 每次处理不同n时，仅需执行“取全局位置→排序位置→重排子集”三步，核心开销是对整数数组的排序，效率远高于重复对x[n]排序，也不需要np.isin这类慢集合操作，非常适合决策树多层循环的场景

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