使用np.argsort()对数组某一列排序

a = np.array([[3,4,3],[2,1,7],[2,3,4],[5,6,2]])

a =
array([[3 4 3]
[2 1 7]
[2 3 4]
[5 6 2]])

np.argsort(a[:,2])

array([3, 0, 2, 1], dtype=int64)

a[[0,1],:]

array([[3, 4, 3],
[2, 1, 7]])

a[np.argsort(a[:,2]),:]

array([[5, 6, 2],
[3, 4, 3],
[2, 3, 4],
[2, 1, 7]])
实现对数组第二列进行排序 列数： 0 1 2

Leetcode 973. 最接近原点的 K 个点

class Solution:
    def kClosest(self, points: List[List[int]], K: int) -> List[List[int]]:
        ### 构建距离信息数组
        import numpy as np
        distance = []
        for i in range(len(points)):
            distance.append([i, points[i][0]**2+points[i][1]**2])
        a = np.array(distance)
        a = a[np.argsort(a[:,1]),:]
        #print(a)
        res = []
        for i in range(K):
            res.append(points[a[i][0]])
        return res