C++で汚染された二分木を復元して要素を検索する方法

問題の概要

次のようなルールに従う二分木を考えます。

• root.val == 0 である
• treeNode.val が x であり、treeNode.left が NULL でない場合、treeNode.left.val = 2 * x + 1 となる
• treeNode.val が x であり、treeNode.right が NULL でない場合、treeNode.right.val = 2 * x + 2 となる

ここで、この二分木は「汚染」されているものとします。つまり、すべてのノードの値が -1 に書き換えられている状態です。まず二分木を復元した上で、以下の FindElements クラスを実装する必要があります。

• FindElements(TreeNode* root)：汚染された二分木を受け取ってオブジェクトを初期化します。最初に木を復元しなければなりません。
• bool find(int target)：復元後の二分木に target の値が存在するかどうかを返します。

たとえば、次のような木が与えられたとします。

復元後に 1、3、5 を検索すると、結果はそれぞれ true、true、false となります。

解法のアプローチ

この問題は、DFS（深さ優先探索）を使って木を一度だけ走査し、復元した値をすべて集合に記録しておくことで効率的に解けます。手順は以下の通りです。

• 整数の集合（set）a を定義します。
• dfs() メソッドを定義します。引数は root と rootVal で、rootVal の初期値は -1 です。
• root が NULL の場合はそのまま戻ります。
• rootVal が -1 の場合はノードの値を 0 に設定し、それ以外の場合は rootVal をそのまま設定します。
• ノードの値を集合 a に挿入します。
• dfs(左の子, 2 * ノードの値 + 1) と dfs(右の子, 2 * ノードの値 + 2) を再帰的に呼び出します。
• 初期化（復元）時には dfs(root, -1) を呼び出します。
• find() では、目的の値が集合 a に存在するかを確認し、存在すれば true、存在しなければ false を返します。

この方法では、復元に O(n) の計算量がかかる一方、find() は std::set（平衡二分探索木）による検索のため O(log n) で高速に判定できます。

以下の実装を見ると、より理解が深まるでしょう。

コード例

#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 insert(TreeNode **root, int val){
   queue<TreeNode*> q;
   q.push(*root);
   while(q.size()){
      TreeNode *temp = q.front();
      q.pop();
      if(!temp->left){
         if(val != NULL)
            temp->left = new TreeNode(val);
         else
            temp->left = new TreeNode(0);
         return;
      }else{
         q.push(temp->left);
      }
      if(!temp->right){
         if(val != NULL)
            temp->right = new TreeNode(val);
         else
            temp->right = new TreeNode(0);
         return;
      }else{
         q.push(temp->right);
      }
   }
}
TreeNode *make_tree(vector<int> v){
   TreeNode *root = new TreeNode(v[0]);
   for(int i = 1; i<v.size(); i++){
      insert(&root, v[i]);
   }
   return root;
}
class FindElements {
   public:
   set <int> a;
   void dfs(TreeNode* root,int rootVal=-1){
      if(!root)return;
      root->val = rootVal == -1?0:rootVal;
      a.insert(root->val);
      dfs(root->left,2*root->val + 1);
      dfs(root->right, 2*root->val + 2);
   }
   FindElements(TreeNode* root) {
      dfs(root);
   }
   bool find(int t) {
      return a.find(t)!=a.end();
   }
};
main(){
   vector<int> v = {-1,-1,-1,-1,-1};
   TreeNode *root = make_tree(v);
   FindElements ob(root);
   cout << (ob.find(1)) << endl;
   cout << (ob.find(3)) << endl;
   cout << (ob.find(5)) << endl;
}

入力

木を [-1,-1,-1,-1,-1] で初期化し、find(1)、find(3)、find(5) を呼び出す

出力

1
1
0

出力は 1（true）、1（true）、0（false）となり、復元後の木には 1 と 3 は存在するが 5 は存在しないことが確認できます。