如何在Python中获取二维数组前n个argmins对应的索引坐标

依赖说明 #

需要用到numpy库做数值计算，如果未安装可以执行pip install numpy安装。

函数实现 #

通用实现（适合小规模数组） #

直接对数组全排序取前n个，逻辑简单易懂：

import numpy as np

def n_argmins(arr, n):
    # 将二维数组扁平化后排序，取前n个最小值的一维索引
    flat_indices = arr.flatten().argsort()[:n]
    # 把一维索引转换为原数组的二维坐标
    return np.array(np.unravel_index(flat_indices, arr.shape)).T

高性能实现（适合大规模数组） #

用分区操作替代全排序，时间复杂度更低，处理大数组时性能优势明显：

def n_argmins_fast(arr, n):
    flat_arr = arr.flatten()
    # 分区操作，仅保证前n个元素是整体最小的n个，无需全排序
    flat_indices = np.argpartition(flat_arr, n-1)[:n]
    # 对前n个元素的索引按值排序，保证输出顺序是从小到大
    flat_indices = flat_indices[np.argsort(flat_arr[flat_indices])]
    return np.array(np.unravel_index(flat_indices, arr.shape)).T

测试验证 #

用你给出的示例测试：

a = np.array([[1, 5, 9], [2, 3, 10], [4, 11, 12]])
argmins = n_argmins(a, 5)
print(argmins)

输出结果（按元素值从小到大排列坐标）：

[[0 0]
 [1 0]
 [1 1]
 [2 0]
 [0 1]]

如果需要调整坐标的排列规则（比如按行、列优先排序，或者调换行列顺序），可以直接对返回的坐标数组做二次排序调整。

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