Numpy 分区比小数组的排序慢

Question

我正在寻找一种有效的方法来计算 numpy 数组中的第 n 个最大值，this answer 将我引导至np.partition。

顺便说一句，我注意到对于少于 100 个条目的数组，朴素排序比 np.partition 方法更快。 （相反，对于大型数组，增益是显而易见的）

小数组的 np.partition 运行时间几乎持平的原因是什么？

生成图片的代码：

import pandas as pd
import numpy as np

import timeit

def func_1(inp):
    return np.partition(inp, 10)[10]

def func_2(inp):
    return np.sort(inp)[10]

a = []
b = []

N_tests = int(1e5)

for wdw in range(20, 1000, 10):

    print wdw

    res1 = timeit.timeit("func_1(test)",
                      setup = "import pandas as pd; import numpy as np; wdw_size = %d; test = np.random.randn(wdw_size); from __main__ import func_1"%wdw, number = N_tests)

    a.append(res1)

    res2 = timeit.timeit("func_2(test)",
                      setup = "import pandas as pd; import numpy as np; wdw_size = %d; test = np.random.randn(wdw_size); from __main__ import func_2"%wdw, number = N_tests)

    b.append(res2)

import matplotlib.pyplot as plt
plt.plot(range(20,1000, 10), a, range(20, 1000, 10), b)
plt.legend(['np.partition', 'np.sort'])
plt.xlabel('Array Size')
plt.ylabel('Time')

Answer 1

根据文档，np.partition 是通过 Introselect 实现的 - 一种最坏情况下性能为 O(n) 的算法。

一句话，Introselect 是快速排序的增强版，在median of medians 的帮助下。

另一方面，np.sort 是使用普通的旧 快速排序 实现的，其最坏情况下的性能为 O(n^2)。 p>

所以比较两者，而np.sort 仅使用快速排序并且可能以 O(n^2) 作为其最坏情况的性能，np.partition 可以通过以下方式避免这种情况必要时退回中位数的中位数以始终确保 O(n)。

---

不完全确定，但np.sort 可能对小型数组更快，因为np.partition 的算法更复杂，开销更大。