无法理解numpy argpartition输出

Question

我正在尝试使用numpy中的arpgpartition，但似乎出了点问题，我似乎无法弄清楚。这是正在发生的事情：

这些是排序数组的前5个元素 norms

np.sort(norms)[:5]
array([ 53.64759445,  54.91434479,  60.11617279,  64.09630585,  64.75318909], dtype=float32)

但是当我使用 indices_sorted = np.argpartition(norms, 5)[:5]

norms[indices_sorted]
array([ 60.11617279,  64.09630585,  53.64759445,  54.91434479,  64.75318909], dtype=float32)

当我认为应该得到与排序数组相同的结果时？

当我使用3作为参数时，它可以正常工作 indices_sorted = np.argpartition(norms, 3)[:3]

norms[indices_sorted]
array([ 53.64759445,  54.91434479,  60.11617279], dtype=float32)

这对我来说没有多大意义，希望有人可以提供一些见解？

编辑：重新解释这个问题，因为argpartition是否保留k个分区元素的顺序更有意义。

Answer 1

我们需要使用按排序顺序保留的索引列表，而不是将第k个参数作为标量输入。因此，要保持第一个5元素的排序性质np.argpartition(a,5)[:5]，只需执行-

np.argpartition(a,range(5))[:5]

这里是一个示例，目的是为了使事情更清楚：

In [84]: a = np.random.rand(10)

In [85]: a
Out[85]: 
array([ 0.85017222,  0.19406266,  0.7879974 ,  0.40444978,  0.46057793,
        0.51428578,  0.03419694,  0.47708   ,  0.73924536,  0.14437159])

In [86]: a[np.argpartition(a,5)[:5]]
Out[86]: array([ 0.19406266,  0.14437159,  0.03419694,  0.40444978,  0.46057793])

In [87]: a[np.argpartition(a,range(5))[:5]]
Out[87]: array([ 0.03419694,  0.14437159,  0.19406266,  0.40444978,  0.46057793])

请注意，这argpartition在性能方面是有意义的，如果我们希望获取一小部分元素的排序索引，那么可以说元素k数占元素总数的一小部分。

让我们使用一个更大的数据集，并尝试获取所有元素的排序索引，以使上述要点更清楚：

In [51]: a = np.random.rand(10000)*100

In [52]: %timeit np.argpartition(a,range(a.size-1))[:5]
10 loops, best of 3: 105 ms per loop

In [53]: %timeit a.argsort()
1000 loops, best of 3: 893 µs per loop

因此，要对所有元素进行排序np.argpartition并不是要走的路。

现在，假设我想只获取具有该大数据集的前5个元素的排序索引，并保持这些元素的顺序-

In [68]: a = np.random.rand(10000)*100

In [69]: np.argpartition(a,range(5))[:5]
Out[69]: array([1647,  942, 2167, 1371, 2571])

In [70]: a.argsort()[:5]
Out[70]: array([1647,  942, 2167, 1371, 2571])

In [71]: %timeit np.argpartition(a,range(5))[:5]
10000 loops, best of 3: 112 µs per loop

In [72]: %timeit a.argsort()[:5]
1000 loops, best of 3: 888 µs per loop

在这里非常有用！