【C++】最大ヒープからK番目に大きい要素を優先度付きキューで効率よく求める方法

このチュートリアルでは、最大ヒープ（max-heap）から K番目に大きい要素 を取り出すプログラムをC++で実装します。

この問題は、優先度付きキュー（priority queue）を使うことで、ヒープ全体をソートすることなく効率的に解くことができます。アルゴリズムの流れは以下の通りです。

アルゴリズムの手順

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

なぜこの方法が有効か

最大ヒープでは親ノードが必ず子ノード以上の値を持つため、「次に大きい候補」は常に「これまでに取り出したノードの子」の中にしか存在しません。そのため、配列全体をソートする O(n log n) のアプローチと比べ、必要な部分だけを探索するこの手法は O(k log k) で済み、k が小さい場合に非常に効率的です。

実装例

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

#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 findKthGreatestElement(Heap &heap, int k) {
    // ペア（値, インデックス）を格納する優先度付きキュー
    priority_queue<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>{ 44, 42, 35, 33, 31, 19, 27, 10, 26, 14 };
    cout << findKthGreatestElement(heap, 4) << endl;
    return 0;
}

コードのポイント

• 優先度付きキューには pair<int, int> を格納し、値とインデックスをセットで管理しています。こうすることで、pop したノードから子ノードの位置を特定できます。
• 子のインデックスがヒープのサイズ n を超えないよう、範囲チェックを行ってから挿入しています。

実行結果

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

33

サンプルのヒープ {44, 42, 35, 33, ...} の中で、4番目に大きい要素は 33 であるため、正しく動作していることが確認できます。

まとめ

本記事では、優先度付きキューを活用して最大ヒープから K番目に大きい要素を求める方法を解説しました。全要素をソートせずに済むため、k が小さいケースでは特に高速に動作します。チュートリアルについて質問がある場合は、コメント欄でお気軽にお尋ねください。