【C++】最小ヒープからk番目に小さい要素を求める方法

このチュートリアルでは、最小ヒープ(min-heap)からk番目に小さい要素を求めるプログラムをC++で作成します。

この問題は、優先度付きキュー(priority queue)を利用することで効率的に解くことができます。まずは、プログラムを完成させるまでの手順を確認していきましょう。

• 最小ヒープを正しい値で初期化します。
• 優先度付きキューを作成し、最小ヒープの根(ルート)ノードを挿入します。
• k − 1回繰り返すループを作成します。
  • キューから最小の要素を取り出します(pop)。
  • 取り出したノードの左の子と右の子を優先度付きキューに追加します。
• ループ終了後、優先度付きキューの先頭にある要素がk番目に小さい要素になっています。
• その値を返します。

アルゴリズムのポイント

この手法が正しく動作するのは、最小ヒープには「親ノードの値は必ず子ノード以下である」という性質が成り立つためです。各ノードの子だけを順次キューに追加していくことで、ヒープ全体をソートすることなく、小さい方から順にk個の要素だけを取り出すことができます。

計算量はO(k log k)となり、ヒープ全体を並べ替えてからk番目を求める方法(全体でO(n log n))よりも効率的です。

コード例

それでは、実際のコードを見てみましょう。

#include <bits/stdc++.h>
using namespace std;
struct Heap {
    vector<int> elements;
    int n;
    Heap(int i = 0): n(i) {
        elements = vector<int>(n);
    }
};
inline int leftIndex(int i) {
    return 2 * i + 1;
}
inline int rightIndex(int i) {
    return 2 * i + 2;
}
int findKthSmallestElement(Heap &heap, int k) {
    priority_queue<pair<int, int>, vector<pair<int, int>>, greater<pair<int, int>>> queue;
    queue.push(make_pair(heap.elements[0], 0));
    for (int i = 0; i < k - 1; ++i) {
        int node = queue.top().second;
        queue.pop();
        int left = leftIndex(node), right = rightIndex(node);
        if (left < heap.n) {
            queue.push(make_pair(heap.elements[left], left));
        }
        if (right < heap.n) {
            queue.push(make_pair(heap.elements[right], right));
        }
    }
    return queue.top().first;
}
int main() {
    Heap heap(10);
    heap.elements = vector<int>{ 10, 14, 19, 24, 32, 41, 27, 44, 35, 33 };
    cout << findKthSmallestElement(heap, 4) << endl;
    return 0;
}

実行結果

上記のコードを実行すると、次のような出力が得られます。

24

サンプルのヒープ {10, 14, 19, 24, 32, 41, 27, 44, 35, 33} を昇順に並べると「10 → 14 → 19 → 24 → …」となるため、4番目に小さい要素である「24」が出力されます。

まとめ

本チュートリアルでは、優先度付きキューを活用して、最小ヒープからk番目に小さい要素を効率的に取得する方法を学びました。ご不明な点がありましたら、コメント欄でお気軽にお知らせください。