C++で二分木の最大幅を求める方法

二分木が与えられたとき、その木の最大幅を求める関数を定義することを考えます。ここでいう木の幅とは、すべてのレベル（階層）の中で最も広いレベルの幅のことを指します。

二分木は完全二分木と同じ構造を持つものとみなしますが、一部のノードはnull（存在しない）である場合があります。あるレベルの幅とは、そのレベルにおける両端ノード（最左端と最右端の非nullノード）間の長さのことであり、両端ノードの間に存在するnullノードも長さの計算に含まれる点に注意が必要です。

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

この場合、最下層のノード構成は [5, 3, null, 9] となるため、最大幅は4になります。

解法のアプローチ

この問題を解くためには、以下の手順に従います。

• 変数 ans := 1、size := 0 として初期化します。
• (ノード, 値) のペアを格納できる両端キュー（deque）q を定義します。
• q に (root, 1) を挿入します。
• q が空でない間、以下の処理を繰り返します。
  • size := q のサイズとします。
  • (ノード, 値) 型のペア curr を定義します。
  • size が 1 の場合は、(先頭要素のノード, 1) を q に追加し、先頭の要素を削除します。
  • size が 0 になるまで、以下を繰り返します。
    • curr := q の先頭要素とし、q から先頭要素を取り出します。
    • curr ノードの左の子が null でない場合、(左の子ノード, 2 × curr の値) を作成して q に挿入します。
    • curr ノードの右の子が null でない場合、(右の子ノード, 2 × curr の値 + 1) を作成して q に挿入します。
    • q のサイズが 1 より大きい場合、ans := max(ans, q の末尾要素の値 − q の先頭要素の値 + 1) と更新します。
    • size := size − 1 とします。
• 最後に ans を返します。

このアルゴリズムのポイント

この手法の鍵となるのは、各ノードに完全二分木上の位置番号を割り当てることです。ルートを 1 とすると、左の子は「親の値 × 2」、右の子は「親の値 × 2 + 1」となります。これにより、たとえ途中にnullノードがあっても、各レベルの幅を「末尾ノードの番号 − 先頭ノードの番号 + 1」として正確に計算できます。処理には幅優先探索（BFS）を用いるため、計算量は O(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:
   int widthOfBinaryTree(TreeNode* root) {
      int ans = 0;
      deque < pair <TreeNode*, int> > q;
      q.push_back({root,1});
      ans = 1;
      int size;
      while(!q.empty()){
         size = q.size();
         pair <TreeNode*, int> curr;
         if(size == 1){
            q.push_back({q.front().first, 1});
            q.pop_front();
         }
         while(size--){
            curr = q.front();
            q.pop_front();
            if(curr.first->left){
               q.push_back({curr.first->left, 2 * curr.second});
            }
            if(curr.first->right){
               q.push_back({curr.first->right, 2 * curr.second + 1});
            }
         }
         if(q.size() > 1)
            ans = max(ans, q.back().second - q.front().second + 1);
      }
      return ans;
   }
};
main(){
   vector<int> v = {1,3,2,5,3,NULL,9};
   TreeNode *root = make_tree(v);
   Solution ob;
   cout << (ob.widthOfBinaryTree(root));
}

入力

[1,3,2,5,3,null,9]

出力

4