如何从一个numpy数组中获得每个数组的N个最大值

Question

我有一个数组x = [[1, 3, 4, 5], [6, 2, 5, 7]]的numpy数组。我想从numpy数组的每个数组中获得N个最大值：[[5, 4], [7, 6]]。我试过使用np.argpartition(x, -N, axis=0)[-N:]，但它给出了ValueError: kth(=-3) out of bounds (1)。做这件事的有效方法是什么？

Answer 1

您可以根据需要对每一行进行排序和切片：

np.sort(x, axis=1)[:, :2]  # --> [[1 3] [2 5]]    2 minimum in each row
np.sort(x, axis=1)[:, 2:]  # --> [[4 5] [6 7]]    2 maximum in each row