Python实现numpy多维数组指定列最小值对应坐标查找

不需要写for循环，也不需要额外创建独立列表存储距离值，直接用NumPy原生的视图切片+索引方法就能实现，全程在原数组上操作，运行效率远高于Python层面的循环遍历。

函数实现 #

import numpy

def closestPoint(arr: numpy.ndarray) -> numpy.ndarray:
    # 切片取所有行的最后一列（距离列）为原数组视图，不生成新的独立数组
    min_distance_row = arr[:, -1].argmin()
    # 返回最小距离对应行的前两位坐标
    return arr[min_distance_row, :2]

测试验证 #

用给出的示例数组运行：

coords = numpy.array([[3,4,5], [5,12,13], [8,15,17]])
print(closestPoint(coords))

运行输出为[3 4]，完全符合预期结果。

实现细节说明 #

• arr[:, -1]是原数组最后一列的视图，没有复制数据，不会额外占用内存创建独立的距离值列表，满足要求
• argmin()是NumPy内置方法，在C层面完成最小值索引查找，比手动写for循环遍历效率高几个量级
• 最终取坐标时直接按行、列切片取值，没有冗余的中间变量或数据拷贝
• 该写法对任意行数的同结构二维数组通用，只要每行最后一位是距离值就不需要修改代码，如果距离值固定在第3列（索引为2），把切片里的-1替换为2效果完全一致。

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