C++で解くTwo Sum IV ― 二分探索木(BST)が入力の場合

問題概要

二分探索木(BST)とターゲット値が1つ与えられます。このとき、BST内に「2つの要素の和がターゲット値と等しくなる」ような組み合わせが存在するかどうかを判定するのが本問題です。

例えば、次のような木が入力として与えられた場合を考えてみましょう。

この場合、出力は True(真)となります。

解法のアプローチ

この問題は、BSTを中間順(inorder)走査して昇順の配列を作り、その後「双方向ポインタ(two pointer)」を使うことで効率的に解けます。具体的には、以下の手順に従います。

1. 値を格納するための配列 v を定義します。
2. 関数 inorder() を定義します(引数は root)。
   • root が NULL の場合は、何もせず return します。
   • まず root の左部分木に対して inorder() を再帰呼び出しします。
   • 次に、root の値を配列 v に挿入します。
   • 最後に、root の右部分木に対して inorder() を再帰呼び出しします。
3. 関数 findnode() を定義します(引数は k)。
   • n := 配列 v のサイズ、i := 0、j := n - 1 として初期化します。
   • i < j の間、以下を繰り返します。
     t := v[i] + v[j] を計算し、
     ・t が k と等しければ true を返す
     ・t が k より小さければ i を1増やす(和を大きくする)
     ・そうでなければ j を1減らす(和を小さくする)
4. ループが終了しても見つからなければ、false を返します。
5. メイン処理では、次のように実行します。
   ・inorder(root) を呼び出す
   ・配列 v をソートする
   ・findnode(k) の結果を返す

計算量について

BSTの中間順走査は O(n)、ソートも O(n log n)(中間順走査の結果は既に昇順のため、実際は不要になることもあります)、双方向ポインタによる探索は O(n) です。全体の時間計算量は O(n log n)、空間計算量は O(n) となります。

実装例(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 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 Solution {
public:
   vector<int> v;
   void inorder(TreeNode* root){
      if (root == NULL || root->val == 0)
         return;
      inorder(root->left);
      v.push_back(root->val);
      inorder(root->right);
   }
   bool findnode(int k){
      int n = v.size(), i = 0, j = n - 1;
      while (i < j) {
         int t = v[i] + v[j];
         if (t == k)
            return true;
         if (t < k)
            i++;
         else
            j--;
      }
      return false;
   }
   bool findTarget(TreeNode* root, int k){
      inorder(root);
      sort(v.begin(), v.end());
      return findnode(k);
   }
};
main(){
   Solution ob;
   vector<int> v = {5,3,6,2,4,NULL,7};
   TreeNode *root = make_tree(v);
   cout << (ob.findTarget(root, 9));
}

入力

{5,3,6,2,4,NULL,7},9

出力

1

まとめ

上記の例では、木の中に存在する 2 + 7 = 9、あるいは 3 + 6 = 9、5 + 4 = 9 といった組み合わせがあるため、結果として 1(true)が出力されます。BSTをソート済み配列に変換してしまえば、あとは古典的な「Sorted Array上のTwo Sum」問題として双方向ポインタで線形時間で解ける、というのがこのアプローチのポイントです。