【C++】先行順トラバーサルの文字列から二分木を復元するアルゴリズムと実装

問題概要

二分木が与えられ、その根ノードに対して先行順（プレオーダー）の深さ優先探索を実行することを考えます。

この走査では、各ノードを訪問するたびに、まずそのノードの深さDと同じ数だけダッシュ「-」を出力し、その直後にノードの値を表示します。深さがDのノードの直接の子の深さはD+1となり、根ノードの深さは0です。

さらに重要なルールとして、あるノードに子が1つしか存在しない場合、その子は必ず左の子であることが保証されています。この走査の出力文字列Sが与えられたとき、元の二分木を復元し、その根を返すのが本問題です。

例えば、入力が「1-2--3--4-5--6--7」の場合、復元される木は次のような構造になります。

        1
       / \
      2   5
     / \ / \
    3  4 6  7

解法のアプローチ：スタックを使った復元

この問題は、スタックを1本使うことで線形時間で解くことができます。基本的な考え方は、「スタックには根から直前のノードまでの経路上のノードを保持させる」というものです。具体的な手順は以下の通りです。

• スタック st を1つ定義する
• i := 0、n := 文字列Sの長さ とする
• lvl := 0、num := 0 と初期化する
• i < n の間、以下の処理を繰り返す
  • lvl := 0 とし、S[i] が '-' である限り lvl を1ずつ増加させながら i を進める（これでノードの深さが求まる）
  • num := 0 とする
  • i < n かつ S[i] が '-' ではない間、num := num × 10 + (S[i] - '0') を計算して i を進め、ノードの値を読み取る
  • st のサイズが lvl より大きい間、要素をポップする（深すぎる不要なノードを取り除き、親候補をスタックトップに揃える）
  • 値 num を持つ新しいツリーノード temp を生成する
  • st が空でなく、スタックトップの左の子がNULLであれば、トップの左の子を temp に設定する
  • そうでなく st が空でなければ、トップの右の子を temp に設定する
  • temp を st にプッシュする
• 最後に、st のサイズが1より大きい間、要素をポップする
• st が空であればNULLを、そうでなければスタックトップを根として返す

なぜスタックでうまくいくのか

先行順走査では、必ず「親 → 左部分木 → 右部分木」の順でノードが現れます。したがって、新しく現れるノードの親は、直前のノードか、その祖先のいずれかです。スタックのサイズを現在の深さに対応させて管理すれば、深さが浅いノードが現れた時点で不要になった子孫をまとめて取り除くことができ、常に正しい親ノードがスタックトップに残ります。

また、「子が1つだけの場合は必ず左の子」という制約のおかげで、左の子が空いていれば左へ、すでに埋まっていれば右へ接続するだけで、木を一意に復元できます。計算量は文字列を一度走査するだけのO(n)（nは文字列長）、スタックには木の高さ分のノードしか積まれないため、空間計算量もO(H)（Hは木の高さ）で抑えられます。

C++による実装例

#include <bits/stdc++.h>
using namespace std;
class TreeNode{
   public:
   int val;
   TreeNode *left, *right;
   TreeNode(int data){
      val = data;
      left = NULL;
      right = NULL;
   }
};
void inord(TreeNode *root){
   if(root != NULL){
      inord(root->left);
      cout << root->val << " ";
      inord(root->right);
   }
}
class Solution {
   public:
   TreeNode* recoverFromPreorder(string S) {
      stack<TreeNode*> st;
      int i = 0;
      int n = S.size();
      int lvl = 0;
      int num = 0;
      while (i < n) {
         for (lvl = 0; S[i] == '-'; lvl++, i++)
         ;
         num = 0;
         while (i < n && S[i] != '-') {
            num = num * 10 + (S[i] - '0');
            i++;
         }
         while (st.size() > lvl)
         st.pop();
         TreeNode* temp = new TreeNode(num);
         if (!st.empty() && !st.top()->left) {
            st.top()->left = temp;
         }
         else if (!st.empty()) {
            st.top()->right = temp;
         }
         st.push(temp);
      }
      while (st.size() > 1)
      st.pop();
      return st.empty() ? NULL : st.top();
   }
};
main(){
   Solution ob;
   TreeNode *root = ob.recoverFromPreorder("1-2--3--4-5--6--7");
   inord(root);
}

入力

"1-2--3--4-5--6--7"

出力

3 2 4 1 6 5 7

出力は復元した木の中順（インオーダー）走査の結果です。「3 2 4 1 6 5 7」という並びが、上記で示した木構造の中順走査と一致していることを確認できます。