ndarray的每一行中的N个最大值

小编典典

ndarray的每一行中的N个最大值

algorithm

我有一个ndarray，其中每一行都是一个单独的直方图。对于每一行，我希望找到前N个值。

我知道全局前N个值的解决方案（一种在numpy数组中找到最大N个元素的快速方法），但是我看不到如何获取每一行的前N个。

我可以遍历每一行并应用一维解决方案，但是我不应该能够通过numpy广播做到这一点吗？

阅读 514

2020-07-28

共1个答案

小编典典

您可以使用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]])

2020-07-28