python - 是否可以按降序使用 argsort？

Question

考虑以下代码:

avgDists = np.array([1, 8, 6, 9, 4])
ids = avgDists.argsort()[:n]

这给了我 n 最小元素的索引。是否可以按降序使用相同的 argsort 来获取 n 最高元素的索引？

Answer 1

如果对数组求反，最低元素将变为最高元素，反之亦然。因此，n 个最高元素的索引为:

(-avgDists).argsort()[:n]

另一种推理方式，如 comments 中所述, 是观察大元素在 argsort 中的 last 到来。因此，您可以从 argsort 的尾部读取以查找 n 最高元素:

avgDists.argsort()[::-1][:n]

这两种方法的时间复杂度都是 O(n log n)，因为 argsort 调用是这里的主要术语。但是第二种方法有一个很好的优势:它将数组的 O(n) 否定替换为 O(1) 切片。如果您在循环内使用小数组，那么您可能会从避免这种否定中获得一些性能提升，如果您使用的是大型数组，那么您可以节省内存使用量，因为否定会创建整个数组的副本。

请注意，这些方法并不总是给出相同的结果:如果向 argsort 请求稳定的排序实现，例如通过传递关键字参数kind='mergesort'，第一个策略将保持排序稳定性，但第二个策略将破坏稳定性(即相等的项目的位置将被反转)。

时间示例:

使用一个包含 100 个 float 和一个长度为 30 的尾部的小数组， View 方法的速度提高了大约 15%

>>> avgDists = np.random.rand(100)
>>> n = 30
>>> timeit (-avgDists).argsort()[:n]
1.93 µs ± 6.68 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
>>> timeit avgDists.argsort()[::-1][:n]
1.64 µs ± 3.39 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
>>> timeit avgDists.argsort()[-n:][::-1]
1.64 µs ± 3.66 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

对于较大的数组，argsort占优势，没有明显的时序差异

>>> avgDists = np.random.rand(1000)
>>> n = 300
>>> timeit (-avgDists).argsort()[:n]
21.9 µs ± 51.2 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
>>> timeit avgDists.argsort()[::-1][:n]
21.7 µs ± 33.3 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)
>>> timeit avgDists.argsort()[-n:][::-1]
21.9 µs ± 37.1 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

请注意 the comment from nedim下面是不正确的。是否在反转之前或之后截断对效率没有影响，因为这两个操作只是以不同的方式跨越数组的 View ，而不是实际复制数据。