python - 无法理解 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

在这里非常有用!