如何从包含最多 3 个值的 2d numpy 数组中获取列的索引

Question

我有一个数组：

a = np.array([[22,11,44,33,66],

              [22,11,2,1,66],

              [1,11,44,22,4],

              [22,11,88,99,66]])

作为输出，我想要一个包含最大 3 个值的索引的数组作为二维数组。例如对于上面的数组输出将是：

array([[4,2,3],

       [4,0,1],

       [2,3,1],

       [3,2,4]])

Answer 1

要获取k数组的顶部元素，请对其进行分区。由于分区通常为您提供k最低元素，请使用反向索引：k = 3top = np.argpartition(a, -k, axis=1)[:, -k:]如果您需要按降序对索引进行排序，请np.argsort与结果一起使用：rows = np.arange(a.shape[0])[:, None]s = np.argsort(a[rows, top], axis=1)[:, ::-1]top = top[rows, s]rows当您使用top和进行花哨的索引时，有必要确保正确选择所有索引s。每行的索引必须颠倒 ( [:, ::-1]) 以获得升序。