C++で二分木の各レベルの平均値を求めるアルゴリズムと実装方法

問題の概要

空でない二分木（バイナリツリー）が与えられたとき、各レベル（階層）に存在するノードの値の平均を計算し、その結果を配列として返すことを考えます。

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

この木は3つのレベルで構成されており、出力は [3, 14.5, 11] となります。具体的には以下のように計算されます。

• 第1レベル：ノード 3 のみ → 平均は 3
• 第2レベル：ノード 9 と 20 → 平均は (9 + 20) / 2 = 14.5
• 第3レベル：ノード 15 と 7 → 平均は (15 + 7) / 2 = 11

解決のためのアプローチ

この問題は、BFS（幅優先探索）を使うことで効率的に解けます。キューを活用してレベルごとにノードを処理し、各レベルの平均値を求めていきます。手順は以下の通りです。

1. 結果を格納するための配列 result を定義します。
2. ノードを管理するためのキュー q を定義し、ルートノードを挿入します。
3. キューが空になるまで、以下の処理を繰り返します。
   • n に現在のキューのサイズ（＝現在のレベルのノード数）を代入します。
   • 現在のレベルの値を一時的に保持する配列 temp を定義します。
   • n が0になるまで、次の処理を繰り返します。
     • キューの先頭要素を t として取り出します。
     • t の値を temp に追加します。
     • キューから先頭要素を削除します。
     • t の左の子が存在すれば、それをキューに挿入します。
     • t の右の子が存在すれば、それをキューに挿入します。
     • n を1減らします。
   • temp のサイズが1の場合、その値をそのまま result の末尾に追加します。
   • temp のサイズが2以上の場合、合計値を計算し、その平均（合計 ÷ 要素数）を result に追加します。
4. 最後に result を返します。

C++による実装例

理解を深めるために、実際のC++コードを見てみましょう。

#include <bits/stdc++.h>
using namespace std;
void print_vector(vector<auto> v){
    cout << "[";
    for(int i = 0; i<v.size(); i++){
        cout << v[i] << ", ";
    }
    cout << "]"<<endl;
}
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<float> averageOfLevels(TreeNode *root){
        vector<float> result;
        queue<TreeNode*> q;
        q.push(root);
        while (!q.empty()) {
            int n = q.size();
            vector<float> temp;
            while (n) {
                TreeNode* t = q.front();
                temp.push_back(t->val);
                q.pop();
                if (t->left && t->left->val != 0)
                    q.push(t->left);
                if (t->right && t->right->val != 0)
                    q.push(t->right);
                n--;
            }
            if (temp.size() == 1)
                result.push_back(temp[0]);
            else if (temp.size() > 1) {
                double sum = 0;
                for (int i = 0; i < temp.size(); i++) {
                    sum += temp[i];
                }
                result.push_back(sum / temp.size());
            }
        }
        return result;
    }
};
main(){
    Solution ob;
    vector<int> v = {3,9,20,NULL,NULL,15,7};
    TreeNode *root = make_tree(v);
    print_vector(ob.averageOfLevels(root));
}

入力例

{3,9,20,NULL,NULL,15,7}

出力例

[3, 14.5, 11]

まとめ

このアルゴリズムでは、BFSを用いて二分木をレベルごとに走査することで、各レベルのノード値の平均をO(N)の時間計算量で効率よく求められます。ここでNは木に含まれるノードの総数です。キューのサイズを利用してレベルの境界を検出するテクニックは、レベル順トラバーサル（Level Order Traversal）の典型的なパターンであり、さまざまな木構造の問題に応用できるので、ぜひ覚えておきましょう。