python - numpy argsort 可以返回较低的关系索引吗？

Question

我有一个 numpy 数组:

foo = array([3, 1, 4, 0, 1, 0])

我想要前 3 项。呼唤

foo.argsort()[::-1][:3]

返回

array([2, 0, 4])

注意值 foo[1] 和 foo[4] 是相等的，因此 numpy.argsort() 通过返回数组中最后出现的项目的索引；即索引 4。

对于我的应用程序，我希望打破平局返回数组中最先出现的项目的索引(此处为索引 1)。我如何有效地实现它？

Answer 1

简单的这个怎么样？

(-foo).argsort(kind='mergesort')[:3]

为什么会这样:

Args 降序排序(不是 np.argsort 所做的)与 args 升序排序(np.argsort 所做的)相同，相反的值。然后您只需要选择前 3 个排序的索引。现在您所需要的就是确保排序稳定，这意味着在出现平局的情况下，首先保留第一个索引。注意:我认为默认的 kind=quicksort 是稳定的，但从文档看来只有 kind=mergesort 保证是稳定的:(https://docs.scipy.org/doc/numpy/reference/generated/numpy.sort.html)

The various sorting algorithms are characterized by their average speed, worst case performance, work space size, and whether they are stable. A stable sort keeps items with the same key in the same relative order. The three available algorithms have the following properties:

kind speed worst case work space stable

‘quicksort’ 1 O(n^2) 0 no

‘mergesort’ 2 O(n*log(n)) ~n/2 yes

‘heapsort’ 3 O(n*log(n)) 0 no