C++で二分木の最大レベル合計を求めるアルゴリズム

問題概要

二分木（バイナリツリー）のルートノードが与えられたとき、ルートのレベルを1とし、その子ノードをレベル2、さらにその下をレベル3として数えていきます。このとき、「全ノードの値の合計が最大になる最小のレベルX」を求めて返すのが本問題の目的です。

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

この場合の出力は「2」となります。各レベルの合計を計算すると以下のようになります。

• レベル1の合計：1
• レベル2の合計：7 + 0 = 7
• レベル3の合計：7 + (-8) = -1

最大の合計はレベル2の「7」であるため、答えは2になります。

解法のアプローチ

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

1. 初期化： level := 1、sum := ルートrの値、ansLevel := level、ansSum := sum と設定します。
2. キューの準備： キューqを定義し、ルートノードrを挿入します。
3. ループ処理： qが空でない間、以下を繰り返します。
   • capacity := qのサイズを取得します。
   • levelを1増やし、sum := 0にリセットします。
   • capacityが0になるまで、以下を繰り返します。
     • node := qの先頭ノードを取り出し、キューから削除します。
     • 右の子ノードが存在する場合は、sumにその値を加算し、右の子ノードをqに挿入します。
     • 左の子ノードが存在する場合は、sumにその値を加算し、左の子ノードをqに挿入します。
     • capacityを1減らします。
   • ansSum < sum であれば、ansSum := sum、ansLevel := level を更新します。
4. 結果の返却： ansLevelを返します。

C++での実装例

それでは、実際の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:
    int maxLevelSum(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(){
    vector<int> v = {1,7,0,7,-8,NULL,NULL};
    TreeNode *root = make_tree(v);
    Solution ob;
    cout <<ob.maxLevelSum(root);
}

入力例

[1,7,0,7,-8,null,null]

出力例

2

計算量について

このアルゴリズムの時間計算量はO(N)です。Nは二分木内のノード総数であり、各ノードを一度だけ訪問するためです。空間計算量もO(N)となり、BFSに使用するキューが最悪の場合、1レベルあたりのノード数分のメモリを保持する必要があるためです。