C++で完全二分木のノード数を効率的に数える方法

完全二分木のノード数を数える問題

完全二分木（Complete Binary Tree）が与えられたとき、その木に含まれるノードの総数を求めるのがこの問題の目的です。例えば、次のような木があった場合、出力は 6 になります。

すべてのノードを一つずつ訪問して数えれば O(n) で解けますが、完全二分木の性質をうまく利用すると、より少ない計算量でノード数を求めることができます。

解法のアプローチ

ここでは再帰的なアプローチを採用します。鍵となるのは、「ある部分木について左端の高さと右端の高さが一致しているなら、その部分木は完全な満木（パーフェクトバイナリツリー）である」という完全二分木の性質です。満木であれば、ノード数は 2^h − 1 という式で即座に計算できます。

countNodes() メソッドは根（ルート）を引数として受け取り、以下の手順で処理を進めます。

1. 高さを表す変数 hr と hl をそれぞれ 0 で初期化します。
2. 根を指す2つのポインタ l と r を作成します。
3. l が空でない間、hl を1ずつ増やしながら l を左の子へ移動し、左側の高さ hl を求めます。
4. r が空でない間、hr を1ずつ増やしながら r を右の子へ移動し、右側の高さ hr を求めます。
5. hl と hr が等しい場合、この部分木は満木なので (2^hl) − 1 を返します。
6. それ以外の場合は、1 + countNodes(左部分木) + countNodes(右部分木) を返します。

理解を深めるために、以下の実装例を見てみましょう。

実装例

#include <bits/stdc++.h>
using namespace std;
class TreeNode{
   public:
      int val;
      TreeNode *left, *right;
      TreeNode(int data){
         val = data;
         left = 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 fastPow(int base, int power){
      int res = 1;
      while(power > 0){
         if(power & 1) res *= base;
         base *= base;
         power >>= 1;
      }
      return res;
   }
   int countNodes(TreeNode* root) {
      int hr = 0;
      int hl = 0;
      TreeNode* l = root;
      TreeNode* r = root;
      while(l){
         hl++;
         l = l->left;
      }
      while(r){
         r = r->right;
         hr++;
      }
      if(hl == hr) return fastPow(2, hl) - 1;
      return 1 + countNodes(root->left) + countNodes(root->right);
   }
};
main(){
   Solution ob;
   vector<int> v = {1,2,3,4,5,6,7,8,9,10};
   TreeNode *node = make_tree(v);
   cout << (ob.countNodes(node));
}

入力

[1,2,3,4,5,6,7,8,9,10]

出力

10

計算量について

このアルゴリズムでは、各再帰呼び出しのたびに木の高さ分だけポインタを移動するため、1回の処理に O(log n) かかります。しかし、左右の高さが一致した時点でその部分木全体のノード数が一度に確定するため、再帰が深く進むのは片側の経路だけで済みます。結果として、全体の時間計算量は O(log n × log n) となり、全ノードを走査する O(n) よりも効率的です。空間計算量は再帰スタックの深さ分の O(log n) です。