C++で二分木の各ノードに次の右ポインタを設定する（Part II）

問題概要

次のような二分木を考えてみましょう。各ノードは (data, left, right, next) というフィールドを持っており、left は左部分木を、right は右部分木を指します。そして next ポインタは、同じ階層における一つ右隣のノードを指すものです。右隣にノードが存在しない場合は null になります。初期状態ではすべての next ポインタが null に設定されているため、これらのリンクを適切に張る必要があります。

例えば、下図のような木が与えられた場合、次のように変換されます。

解法のアプローチ

この問題は、next ポインタそのものを利用して各レベルを連結リストとして走査することで、BFS 用のキューなどの追加メモリを使わずに解くことができます。現在のレベルを順にたどりながら、次のレベルのノード同士を prev ポインタでつなげていくのがポイントです。

アルゴリズムの手順

• pre := root、nextPre := null、prev := null と初期化する
• pre が null でない限り、以下を繰り返す
  • pre が null でない限り、以下を繰り返す
    • pre の左の子が存在する場合
      • prev が null でなければ、prev の next を pre の左の子に設定する。null の場合は、nextPre を pre の左の子に設定する
      • prev := pre の左の子
    • pre の右の子が存在する場合
      • prev が null でなければ、prev の next を pre の右の子に設定する。null の場合は、nextPre を pre の右の子に設定する
      • prev := pre の右の子
  • pre := nextPre として、次のレベルへ移動する
  • nextPre と prev を null にリセットする
• root を返す

このアルゴリズムの計算量は、時間計算量 O(n)、追加の空間計算量 O(1) であり、非常に効率的です。

C++での実装例

理解を深めるために、以下の実装例を見てみましょう。

#include <bits/stdc++.h>
#include <stack>
using namespace std;
class Node {
   public:
   int val;
   Node* left;
   Node* right;
   Node* next;
   Node() : val(0), left(NULL), right(NULL), next(NULL) {}
   Node(int _val) : val(_val), left(NULL), right(NULL), next(NULL) {}
   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){
            //cout << pre->val << endl;
            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;
         }
         //cout << "*" << endl;
         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->next)
         // q.push(curr->next);
         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);
   Node nodeFive(5);
   Node nodeSeven(7);
   Node nodeTwo(2,&nodeFour,&nodeFive);
   Node nodeThree(3,NULL,&nodeSeven);
   Node nodeOne(1,&nodeTwo,&nodeThree);
   root = &nodeOne;
   Solution ob;
   root = ob.connect(root);
   printTree(root);
}

入力

[1,2,3,4,5,null,7]
Node* root;
Node nodeFour(4);
Node nodeFive(5);
Node nodeSeven(7);
Node nodeTwo(2,&nodeFour,&nodeFive);
Node nodeThree(3,NULL,&nodeSeven);
Node nodeOne(1,&nodeTwo,&nodeThree);
root = &nodeOne;

出力

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