如何从包含最多 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 = 3

top = 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]) 以获得升序。