起因
上個月的一場面試中，面試官問了我這個問題，當時沒能很好地回答出來。後來認真地思考了這個問題，用此文記錄下來。
排序演算法
符合第一直覺的回答應該是快速排序，在系統學過排序演算法後，我們都知道快排是一種很快的排序方法。我們先對數據進行排序，然後再取出其中前 1000 個數。不過你要是真的在面試中只那麼回答的話，那麼基本也無緣後面的面試了。這種思路在輸入的數據規模很小時，不失為一種便捷的方法，但是當數據規模擴大時（題目中的 1 億），其對運行的記憶體要求也將增大。在快排中，其時間複雜度為 O(n log n)，要同時對 1 億整數排序，需要將其讀入記憶體中，以每個 int 類型佔 4 位元組為例，1 億整數則需要佔用 400MB 的記憶體。雖然現代的電腦可用記憶體遠遠大於 400MB，但當數據規模再次擴大到 10 億甚至是 100 億時，其對記憶體的要求則上升到了 4GB 和 40GB，演算法的可用性 (availability) 和擴展性 (scalability) 將受到嚴重的影響。同時，在這種情況下還需要做全排序的話，就只能借助外部排序、位圖排序、基數排序、桶排序 等演算法來解決記憶體不足的問題了。同時，題目中只要求取 top 1000 的數據，對所有元素排序很明顯是不必要的。
我寫一個簡單的演示程式，程式碼如下：
#include <iostream>
#include <algorithm>
#include <random>
#include <chrono>
using namespace std;
#define MAX INT32_MAX
#define MIN INT32_MIN
#define SIZE 100000000 // 100 million
#define K 1000
int main() {
auto start = chrono::high_resolution_clock::now();
// Create a random number generator
random_device rd;
mt19937 gen(rd());
uniform_int_distribution<> dis(MIN, MAX);
// Create an array of random numbers
int *arr = new int[SIZE];
for (int i = 0; i < SIZE; i++) {
arr[i] = dis(gen);
}
auto end = chrono::high_resolution_clock::now();
chrono::duration<double> diff = end - start;
cout << "Time taken to create array: " << diff.count() << " s" << endl;
start = chrono::high_resolution_clock::now();
// sorting via std::sort() - optimized quicksort
sort(arr, arr + SIZE);
end = chrono::high_resolution_clock::now();
diff = end - start;
// the arr[0]~arr[k-1] is the top k elements in arr.
cout << "Time taken to get top K elements: " << diff.count() << " s" << endl;
free(arr);
return 0;
}
# CPU: i7-8700k
# Memory: 16 GB
# OS: Ubuntu 22.10 Kinetic
Time taken to create array: 2.13481 s
Time taken to get top K elements: 22.3849 s
局部淘汰法
為了避免上面的演算法中的不必要排序，我們可以選取出待排序陣列中的前 k (= 1000) 個元素構建一個最小堆。然後從 k + 1 開始與容器中的最小元素 m 相比，如果大於 m，則替換 m 並重新調整最小堆。
其演算法如下：
- 選取前 K 個元素，構建一個最小堆。
- 從第 K+1 個元素開始，與堆頂元素比較，如果比堆頂元素大，則替換堆頂元素，並調整堆。
- 重複步驟 2，直到陣列遍歷完畢。
- 堆中的元素即為前 K 個最大元素。
其程式碼如下：
#include <iostream>
#include <algorithm>
#include <random>
#include <chrono>
using namespace std;
#define MAX INT32_MAX
#define MIN INT32_MIN
#define SIZE 100000000 // 100 million
#define K 1000
int main() {
// Create a random number generator
auto start = chrono::high_resolution_clock::now();
random_device rd;
mt19937 gen(rd());
uniform_int_distribution<> dis(MIN, MAX);
// Create an array of random numbers
int *arr = new int[SIZE];
for (int i = 0; i < SIZE; i++) {
arr[i] = dis(gen);
}
auto end = chrono::high_resolution_clock::now();
chrono::duration<double> diff = end - start;
cout << "Time taken to create array: " << diff.count() << " s" << endl;
// clock start
start = chrono::high_resolution_clock::now();
// generate a min-heap with first k elements
int *heap = new int[K];
for (int i = 0; i < K; i++) {
heap[i] = arr[i];
}
make_heap(heap, heap + K, greater<>());
for (int i = K; i < SIZE; i++) {
if (arr[i] > heap[0]) {
pop_heap(heap, heap + K, greater<>());
heap[K - 1] = arr[i];
push_heap(heap, heap + K, greater<>());
}
}
// clock end
end = chrono::high_resolution_clock::now();
diff = end - start;
cout << "Time taken to get top K elements: " << diff.count() << " s" << endl;
free(arr);
free(heap);
return 0;
}
Time taken to create array: 2.18854 s
Time taken to get top K elements: 0.178133 s
C++ STL set
我們知道，標準庫中的 set 底層實作是一棵紅黑樹（平衡二元樹），在這裡我想到了是不是可以借助 set 容器來取 top k 元素呢。如果允許重複元素的話，則使用 multiset 來代替 set.
程式碼如下：
#include <iostream>
#include <random>
#include <chrono>
#include <set>
using namespace std;
#define MAX INT32_MAX
#define MIN INT32_MIN
#define SIZE 100000000 // 100 million
#define K 1000
int main() {
// Create a random number generator
auto start = chrono::high_resolution_clock::now();
random_device rd;
mt19937 gen(rd());
uniform_int_distribution<> dis(MIN, MAX);
// Create an array of random numbers
int *arr = new int[SIZE];
for (int i = 0; i < SIZE; i++) {
arr[i] = dis(gen);
}
auto end = chrono::high_resolution_clock::now();
chrono::duration<double> diff = end - start;
cout << "Time taken to create array: " << diff.count() << " s" << endl;
// clock start
start = chrono::high_resolution_clock::now();
// Create a set of K the largest numbers
set<int> s;
for (int i = 0; i < SIZE; i++) {
if (s.size() < K) {
s.insert(arr[i]);
} else {
if (arr[i] > *s.begin()) {
s.erase(s.begin());
s.insert(arr[i]);
}
}
}
// clock end
end = chrono::high_resolution_clock::now();
diff = end - start;
cout << "Time taken to get top K elements: " << diff.count() << " s" << endl;
// Print the set in reverse order
// for (auto it = s.rbegin(); it != s.rend(); it++) {
// cout << *it << endl;
// }
free(arr);
return 0;
}
Time taken to create array: 2.16181 s
Time taken to get top K elements: 1.68677 s
總結
透過上面的實驗，我們可以看到，使用局部淘汰的思想和最小堆的資料結構來取 top k 元素的效率是最高的，其次是使用 set 容器，最後是使用排序演算法。