您可以使用np.partition
与链接的问题相同的方式:排序已经沿最后一个轴进行:
In [2]: a = np.array([[ 5, 4, 3, 2, 1],
[10, 9, 8, 7, 6]])
In [3]: b = np.partition(a, -3) # top 3 values from each row
In [4]: b[:,-3:]
Out[4]:
array([[ 3, 4, 5],
[ 8, 9, 10]])