C++でノード値の合計が最小となる二分木のレベルを求めるプログラム

二分木（バイナリツリー）を考えます。根（ルート）のレベルを1とし、その子のレベルを2、さらにその下のレベルを3というように定義します。このとき、レベルXに存在するすべてのノードの値の合計が最小になるような、最も小さいレベルXを見つけるのが本記事の目的です。

例として、次のような二分木を考えてみましょう。

この場合、出力は 2 となります。なぜなら、レベル2のノードの値の合計は 4 + (-10) = -6 となり、これが全レベルの中で最小だからです。

解法のアプローチ

この問題は、幅優先探索（BFS）を使って各レベルごとにノードの値の合計を計算し、その中で最小となるレベルを記録していくことで解けます。具体的には、以下の手順に従います。

• level := 1、sum := ルートrの値、ansLevel := level、ansSum := sum で初期化する。
• キュー q を定義し、ルートノード r を挿入する。
• キュー q が空でない間、以下を繰り返す。
  • capacity := キューqのサイズ を取得する。
  • level を1増やし、sum := 0 にリセットする。
  • capacity が0になるまで以下を繰り返す。
    • キューの先頭ノードを取り出し、キューから削除する。
    • 右の子が存在する場合、sum に右の子の値を加算し、右の子をキューに挿入する。
    • 左の子が存在する場合、sum に左の子の値を加算し、左の子をキューに挿入する。
    • capacity を1減らす。
  • ansSum > sum の場合、ansSum := sum、ansLevel := level と更新する。
• ansLevel を返す。

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;
      }
};
class Solution {
   public:
   int solve(TreeNode* r) {
      int level = 1, sum = r->val;
      int ansLevel = level, ansSum = sum;
      queue <TreeNode*> q;
      q.push(r);
      while(!q.empty()){
         int capacity = q.size();
         level++;
         sum = 0;
         while(capacity--){
            TreeNode* node = q.front();
            q.pop();
            if(node->right){
               sum += node->right->val;
               q.push(node->right);
            }
            if(node->left){
               sum += node->left->val;
               q.push(node->left);
            }
         }
         if(ansSum>sum){
            ansSum = sum;
            ansLevel = level;
         }
      }
      return ansLevel;
   }
};
main(){
   TreeNode *root = new TreeNode(5);
   root->left = new TreeNode(4);
   root->right = new TreeNode(-10);
   root->left->right = new TreeNode(-2);
   root->right->left = new TreeNode(-7);
   root->right->right = new TreeNode(15);
   Solution ob;
   cout <<ob.solve(root);
}

入力

TreeNode *root = new TreeNode(5);
root->left = new TreeNode(4);
root->right = new TreeNode(-10);
root->left->right = new TreeNode(-2);
root->right->left = new TreeNode(-7);
root->right->right = new TreeNode(15);

出力

2

アルゴリズムのポイント

この実装では、キューの現在のサイズを capacity として保存することで、「同じレベルのノードだけを処理する」というBFSの重要な性質を実現しています。各レベルの処理が終わった時点で sum にはそのレベルのノード値の合計が入るため、これまでの最小値 ansSum と比較して、より小さい場合は ansLevel を更新します。

計算量についても触れておきましょう。すべてのノードをちょうど1回ずつ訪問するため、時間計算量はO(n)（nはノード数）、キューに格納されるノード数は最大で木の幅程度なので、空間計算量もO(n) となります。