C++で解く二分木における最大合計BSTの求め方

問題概要

二分木のルートが与えられたとき、二分探索木（BST）でもある部分木の中から、ノード値の合計が最大となるものを見つけることを考えます。

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

この場合の出力は 20 になります。これは選択されたBSTに含まれるすべてのノードの値の合計です。

解法のアプローチ

この問題は、木を後順トラバーサル（子→親の順）で処理し、各ノードを根とする部分木に関する情報をボトムアップに集約することで効率的に解けます。具体的な手順は以下の通りです。

• Data という構造体を作成します。sz（部分木のノード数）、maxVal（最大値）、minVal（最小値）、ok（BSTであるかの判定フラグ）、sum（ノード値の合計）というメンバーを持たせます。また、(sz, minVal, maxVal, ok) の順で値を受け取り、sum を 0 に初期化するコンストラクタも定義しておきます。

• 答えを保持する変数 ret を 0 で初期化します。

• solve() というメソッドを定義します。引数には木のノードを渡します。

• ノードが存在しない、またはノードの値が 0 である場合は、Data(0, inf, -inf, true) で初期化した新しい Data オブジェクトを返します（空の部分木は常にBSTとみなせるためです）。

• left := solve(node の左の子)、right := solve(node の右の子) として、左右の部分木の結果を再帰的に取得します。

• Data 型のインスタンス curr を作成し、curr.ok := false で初期化します。

• node->val ≥ right.minVal の場合は、BSTの条件が崩れるため curr を返します。

• node->val ≤ left.maxVal の場合も同様に curr を返します。

• left.ok と right.ok がどちらも真（左右の部分木ともBST）である場合は、以下のように情報を更新します。

  • curr.sum := node->val + left.sum + right.sum

  • ret := max(curr.sum, ret)

  • curr.sz := 1 + left.sz + right.sz

  • curr.ok := true

  • curr.maxVal := max(node->val, right.maxVal)

  • curr.minVal := min(node->val, left.minVal)

• 最後に curr を返します。

• メインメソッドでは、ret を 0 にリセットし、solve(root) を呼び出したうえで ret を返します。

計算量

各ノードを一度だけ訪問するため、時間計算量は O(N)（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;
}
struct Data{
   int sz;
   int maxVal;
   int minVal;
   bool ok;
   int sum;
   Data(){}
   Data(int a, int b, int c, bool d){
      sz = a;
      minVal = b;
      maxVal = c;
      ok = d;
      sum = 0;
   }
};
class Solution {
   public:
   int ret = 0;
   Data solve(TreeNode* node){
      if (!node || node->val == 0)
      return Data(0, INT_MAX, INT_MIN, true);
      Data left = solve(node->left);
      Data right = solve(node->right);
      Data curr;
      curr.ok = false;
      if (node->val >= right.minVal) {
         return curr;
      }
      if (node->val <= left.maxVal) {
         return curr;
      }
      if (left.ok && right.ok) {
         curr.sum = node->val + left.sum + right.sum;
         ret = max(curr.sum, ret);
         curr.sz = 1 + left.sz + right.sz;
         curr.ok = true;
         curr.maxVal = max(node->val, right.maxVal);
         curr.minVal = min(node->val, left.minVal);
      }
      return curr;
   }
   int maxSumBST(TreeNode* root){
      ret = 0;
      solve(root);
      return ret;
   }
};
main(){
   Solution ob;
   vector<int> v =
   {1,4,3,2,4,2,5,NULL,NULL,NULL,NULL,NULL,NULL,4,6};
   TreeNode *root = make_tree(v);
   cout << (ob.maxSumBST(root));
}

入力

{1,4,3,2,4,2,5,NULL,NULL,NULL,NULL,NULL,NULL,4,6}

出力

20

この例では、値 3 を根とする部分木（2, 3, 4, 5, 6）がBSTの条件を満たしており、その合計は 2 + 3 + 4 + 5 + 6 = 20 となります。