有 numpy 有限的 argsort 吗？ [复制]答案

【问题标题】：Is there a numpy limited argsort? [duplicate]有 numpy 有限的 argsort 吗？ [复制]
【发布时间】：2021-10-07 05:58:01
【问题描述】：

我有一个包含 1M 浮点数的数组，想要获得前 10 名。想知道是否有一个 argsort 只会给出前 10 名来加快速度。我想我可以执行 np.argmax() 10 次，每次都删除结果索引（可能通过使其 -inf，不必跟踪移动索引）。只是检查是否有众所周知的内置方式。

不一定是 numpy。其次是另一个非常常见的库。

【问题讨论】：

为什么不直接使用 np.argsort()？
@sehan2 因为我不想计算整个排序。我认为获得前 10 名会快很多
在 CS 中这被称为部分排序（这应该有助于您使用谷歌搜索）。您也许可以通过某种方式利用numpy.partition()。
numpy.argpartition 是我在 numpy 中所知道的最接近的东西。我相信您可以通过arr.argpartition(kth=(0,1,2,3,4,5,6,7,8,9))[:10] 获取数组arr 中前10 个元素（按其排序顺序）的索引。
@AriCooper-Davis 够公平的。虽然额外的处理就像用初始列表中的 (element, index) 对替换每个元素一样简单。

标签： python numpy sorting

【解决方案1】：

你确实可以使用numpy.argpartition():

import numpy
a = numpy.random.rand(1000000)
b = numpy.argpartition(a, [0,1,2,3,4,5,6,7,8,9])[:10]

它快了将近 10 倍（在本例中，在我的机器上）：

>>> import timeit, numpy
>>> a = numpy.random.rand(1000000)
>>> # Sorting traditionally
>>> timeit.timeit('numpy.argsort(a)[:10]', globals=globals(), number=100)
10.2261...
>>> # Getting the top 10 but not in order
>>> timeit.timeit('numpy.argpartition(a, 10)[:10]', globals=globals(), number=100)
1.2125...
>>> # Getting the top 10 in order
>>> timeit.timeit('numpy.argpartition(a, [0,1,2,3,4,5,6,7,8,9])[:10]', globals=globals(), number=100)
1.4764...

【讨论】：