C++で二分木の各レベルのノードをソートして出力する方法

この問題では、二分木が与えられ、各レベルに存在するすべてのノードを値の順序(ソート済み)で出力することが求められます。

まず、具体例を見ながら概念を理解していきましょう。

入力 −

出力 −

20
6 15
2 17 32 78

解決のアプローチ

この問題を解くには、木の各レベルごとにノードの値をソートした状態で出力する必要があります。そのために、以下のデータ構造を利用します。

• queue(キュー):幅優先探索(BFS)のようにノードをたどるために使用
• priority_queue × 2つ:1つは「現在のレベル」の値を昇順で保持し、もう1つは「次のレベル」の値を一時的に保持するために使用(greater<int> を指定することで最小ヒープとして動作)
• NULLセパレータ:キュー内でレベルとレベルの境界を区切るマーカーとして使用

処理の流れ

1. ルートノードとNULLセパレータをキューに追加し、ルートの値を current_level に格納します。
2. キューからノードを取り出しながら、current_level の先頭(最小値)から順に出力します。
3. 子ノードが存在する場合は、キューに子を追加すると同時に、その値を next_level に格納します。
4. NULLセパレータに到達したら改行を出力し、current_level と next_level を swap() で入れ替えて次のレベルへ移行します。
5. キューが空になるまでこの処理を繰り返します。

実装例

上記のロジックを示すC++プログラムは以下の通りです。

#include <iostream>
#include <queue>
#include <vector>
using namespace std;
struct Node {
    int data;
    struct Node *left, *right;
};
void printLevelElements(Node* root){
    if (root == NULL)
        return;
    queue<Node*> q;
    priority_queue<int, vector<int>, greater<int> > current_level;
    priority_queue<int, vector<int>, greater<int> > next_level;
    q.push(root);
    q.push(NULL);
    current_level.push(root->data);
    while (q.empty() == false) {
        int data = current_level.top();
        Node* node = q.front();
        if (node == NULL) {
            q.pop();
            if (q.empty())
                break;
            q.push(NULL);
            cout << "\n";
            current_level.swap(next_level);
            continue;
        }
        cout << data << " ";
        q.pop();
        current_level.pop();
        if (node->left != NULL) {
            q.push(node->left);
            next_level.push(node->left->data);
        }
        if (node->right != NULL) {
            q.push(node->right);
            next_level.push(node->right->data);
        }
    }
}
Node* insertNode(int data){
    Node* temp = new Node;
    temp->data = data;
    temp->left = temp->right = NULL;
    return temp;
}
int main(){
    Node* root = insertNode(12);
    root->left = insertNode(98);
    root->right = insertNode(34);
    root->left->left = insertNode(76);
    root->left->right = insertNode(5);
    root->right->left = insertNode(12);
    root->right->right = insertNode(45);
    cout << "Elements at each Level of binary tree are \n";
    printLevelElements(root);
    return 0;
}

実行結果

Elements at each Level of binary tree are
12
34 98
5 12 45 76

計算量について

すべてのノードを一度ずつ処理するため、基本の計算量は O(n) ですが、各ノードの値を優先度付きキューに挿入・取り出す際のヒープ操作が O(log n) かかるため、全体の時間計算量は O(n log n) となります。空間計算量については、キューと優先度付きキューに最大で O(n) 分の要素が保持される可能性があります。