如何获取大型numpy多维数组中x个最小元素的索引（支持任意维度）

对于大型Numpy多维数组，最快的方法是结合np.argpartition和np.unravel_index——前者通过部分排序避免全量排序的性能开销，后者负责将一维索引还原为原数组的多维索引，完全支持任意维度。

具体步骤： #

• 第一步：使用np.argpartition获取展平数组中前x个最小元素的一维索引。argpartition的时间复杂度为O(n)，远快于全排序的O(n log n)，处理大型数组时优势显著。
• 第二步：通过np.unravel_index将一维索引转换为原多维数组的索引结构。

代码示例： #

假设我们有一个3维数组，要提取其中5个最小元素的索引：

import numpy as np

# 生成测试用的大型多维数组
arr = np.random.rand(100, 100, 100)
x = 5

# 获取展平后前x个最小元素的一维索引
flat_indices = np.argpartition(arr.flatten(), x)[:x]

# 转换为原多维数组的索引
multi_indices = np.unravel_index(flat_indices, arr.shape)

# 输出结果，multi_indices是元组，每个元素对应原数组一个维度的索引
print("x个最小元素的多维索引：", multi_indices)

补充说明： #

• 如果需要这x个元素按从小到大的顺序排列，可以对获取到的一维索引对应的元素值排序，再重新提取索引：

# 获取前x个最小元素的值和对应一维索引
values = arr.flatten()[flat_indices]
sorted_indices = flat_indices[np.argsort(values)]
# 转换为排序后的多维索引
sorted_multi_indices = np.unravel_index(sorted_indices, arr.shape)

这种方式依然比直接对整个数组做argsort高效，因为仅对x个元素做排序，当x远小于数组总元素数时优势明显。

• 该方法完全适配任意维度的数组，不管是2维、3维还是更高维度，np.unravel_index都会根据原数组的shape自动完成索引转换。

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