C++で完全二分木の各ノードにnextポインタ(次の右ポインタ)を設定する方法

問題概要

完全二分木を考えます。各ノードは (data, left, right, next) という4つのフィールドを持っており、left は左部分木を、right は右部分木を指します。next ポインタは、同じレベル(階層)における「次のノード」を指すためのもので、右隣にノードが存在しない場合は null となります。

初期状態ではすべての next ポインタが null に設定されているため、これらのリンクを適切に張り直すことが本記事の目的です。例えば、以下のような木がある場合、

これを次のように変換します。

解法のアプローチ

この問題は、レベルごとにノードをたどりながら next ポインタを順番に接続していくことで解決できます。具体的には、次の3つのポインタを管理します。

• pre:現在処理中のレベルの走査位置を示すポインタ
• nextPre:次のレベルの先頭ノードを記録するポインタ
• prev:同レベル内で直前に処理した子ノードを記録するポインタ

処理の手順は以下の通りです。

• pre := root、nextPre := null、prev := null と初期化する
• pre が null でない限り、以下を繰り返す
  • pre が null でない限り、以下を繰り返す
    • pre の左の子が存在する場合
      • prev が null でなければ prev->next := pre->left、そうでなければ nextPre := pre->left
      • prev := pre->left
    • pre の右の子が存在する場合
      • prev が null でなければ prev->next := pre->right、そうでなければ nextPre := pre->right
      • prev := pre->right
    • pre := pre->next(同じレベル内で次のノードへ移動)
  • pre := nextPre(次のレベルの先頭へ移動)
  • nextPre と prev を null にリセットする
• root を返す

この手法のポイントは、キューなどの追加データ構造を使わず、すでに構築済みの next ポインタを頼りにレベル順走査を実現している点です。そのため、空間計算量は O(1) に抑えられ、時間計算量は全ノードを一度ずつ訪問するため O(n) となります。

C++での実装例

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

#include <bits/stdc++.h>
#include <stack>
using namespace std;
class Node {
public:
    int val;
    Node* left;
    Node* right;
    Node* next;
    Node() {}
    Node(int _val, Node* _left, Node* _right) {
        val = _val;
        left = _left;
        right = _right;
        next = NULL;
    }
};
class Solution {
public:
    Node* connect(Node* root) {
        Node* pre = root;
        Node* nextPre = NULL;
        Node* prev = NULL;
        while(pre){
            while(pre){
                if(pre->left){
                    if(prev){
                        prev->next = pre->left;
                    }else{
                        nextPre = pre->left;
                    }
                    prev = pre->left;
                }
                if(pre->right){
                    if(prev){
                        prev->next = pre->right;
                    }else{
                        nextPre = pre->right;
                    }
                    prev = pre->right;
                }
                pre = pre->next;
            }
            pre = nextPre;
            nextPre = NULL;
            prev = NULL;
        }
        return root;
    }
};
void printTree(Node* root) {
    cout << "[";
    if (root == NULL) return;
    queue<Node*> q;
    Node *curr;
    q.push(root);
    q.push(NULL);
    while (q.size() > 1) {
        curr = q.front();
        q.pop();
        if (curr == NULL){
            q.push(NULL);
        }
        else {
            if(curr->left)
                q.push(curr->left);
            if(curr->right)
                q.push(curr->right);
            if(curr->val == 0){
                cout << "null" << ", ";
            }else{
                cout << curr->val << ", ";
                if (curr->next == NULL) cout<<"#, ";
            }
        }
    }
    cout << "]"<<endl;
}
int main() {
    Node* root;
    Node nodeFour(4, NULL, NULL);
    Node nodeFive(5, NULL, NULL);
    Node nodeSeven(7, NULL, NULL);
    Node nodeSix(6, NULL, NULL);
    Node nodeTwo(2,&nodeFour,&nodeFive);
    Node nodeThree(3,&nodeSix,&nodeSeven);
    Node nodeOne(1,&nodeTwo,&nodeThree);
    root = &nodeOne;
    Solution ob;
    root = ob.connect(root);
    printTree(root);
}

入力

[1,2,3,4,5,6,7]
Node* root;
Node nodeFour(4, NULL, NULL);
Node nodeFive(5, NULL, NULL);
Node nodeSeven(7, NULL, NULL);
Node nodeSix(6, NULL, NULL);
Node nodeTwo(2,&nodeFour,&nodeFive);
Node nodeThree(3,&nodeSix,&nodeSeven);
Node nodeOne(1,&nodeTwo,&nodeThree);
root = &nodeOne;
Solution ob;
root = ob.connect(root);

出力

[1, #, 2, 3, #, 4, 5, 6, 7, #]

出力の # は、そのノードの next ポインタが null(レベルの終端)であることを意味します。この結果から、ルート「1」の次は null、第2レベルでは「2 → 3」、第3レベルでは「4 → 5 → 6 → 7」と正しく連結されていることが確認できます。

まとめ

本記事では、完全二分木の各ノードの next ポインタを、同じレベルの右隣のノードへ接続する方法を解説しました。既存の next ポインタを活用したレベル順走査により、追加メモリ O(1)・時間計算量 O(n) という効率的な実装が可能です。木構造の問題では、このように「すでに持っている情報を再利用する」発想が重要になるので、ぜひ覚えておきましょう。