きっかけ
先月の面接で、面接官からこの質問をされましたが、うまく答えられませんでした。その後、真剣にこの問題について考えたので、この記事にまとめます。
ソートアルゴリズム
最初に直感的に思いつく回答は「クイックソート」でしょう。ソートアルゴリズムを体系的に学んだ人なら、クイックソートが高速なソート方法であることを知っています。まずデータをソートしてから、その中から上位1000個を取り出します。しかし、面接でそのようにだけ答えてしまうと、その後の面接はほぼ望み薄でしょう。この考え方は、入力データの規模が小さい場合には便利な方法ですが、データ規模が大きくなると（問題の1億）、実行時のメモリ要件も増大します。クイックソートの時間計算量は O(n log n) で、1億個の整数を同時にソートするには、それらをメモリに読み込む必要があります。int型が4バイトを占めるとすると、1億個の整数には400MBのメモリが必要です。現代のコンピュータの使用可能メモリは400MBをはるかに超えていますが、データ規模がさらに10億、100億と拡大すると、メモリ要件は4GB、40GBに跳ね上がり、アルゴリズムの可用性 (availability) と拡張性 (scalability) に深刻な影響を及ぼします。また、このような状況で全ソートを行う必要がある場合は、外部ソート、ビットマップソート、基数ソート、バケットソート などのアルゴリズムを利用してメモリ不足を解決するしかありません。同時に、問題では上位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要素を取得できないかと考えました。重複要素を許可する場合は、setの代わりにmultisetを使用します。
コードは以下の通りです：
#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コンテナ、最後にソートアルゴリズムを使用する方法であることがわかりました。