C++で二分木の葉ノードを繰り返し収集・削除するアルゴリズム

問題の概要

二分木が与えられているとします。まずすべての葉（子ノードを持たないノード）を収集して取り除き、その操作を木が空になるまで繰り返します。

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

このとき、出力は [[4,5,3],[2],[1]] となります。最初のラウンドで葉である 4、5、3 が取り除かれ、続いて 2 が、最後に根の 1 が残るという流れです。

解法のアプローチ

この問題は、各ノードの「高さ」（最も深い葉から数えた距離）をDFSで求めると効率的に解けます。同じ高さを持つノードは、必ず同じラウンドで葉になるためです。具体的な手順は以下の通りです。

• マップ sz を定義する

• 2次元配列 ret を定義する

• 引数にノードを取る関数 dfs() を定義する

• node が null の場合は 0 を返す

• sz[node の値] に「1 + max(dfs(node の左の子), dfs(node の右の子))」を代入する

• ret のサイズが sz[node の値] より小さい場合、空の配列 temp を作成して ret の末尾に追加する

• ret[sz[node の値] - 1] の末尾に node の値を挿入する

• sz[node の値] を返す

メインの処理では、dfs(root) を呼び出した後、結果として ret を返すだけで完成です。

実装例

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

#include <bits/stdc++.h>
using namespace std;
void print_vector(vector<vector<auto>> v){
   cout << "[";
   for(int i = 0; i<v.size(); i++){
      cout << "[";
      for(int j = 0; j <v[i].size(); j++){
         cout << v[i][j] << ", ";
      }
      cout << "],";
   }
   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:
   unordered_map <int, int> sz;
   vector < vector <int> > ret;
   int dfs(TreeNode* node){
      if(!node) return 0;
      sz[node->val] = 1 + max(dfs(node->left), dfs(node->right));
      if(ret.size() < sz[node->val]){
         vector <int> temp;
         ret.push_back(temp);
      }
      ret[sz[node->val] - 1].push_back(node->val);
      return sz[node->val];
   }
   vector<vector<int>> findLeaves(TreeNode* root) {
      dfs(root);
      return ret;
   }
};
main(){
   Solution ob;
   vector<int> v = {1,2,3,4,5};
   TreeNode *root = make_tree(v);
   print_vector(ob.findLeaves(root));
}

入力

{1,2,3,4,5}

出力

[[3, 5, 4],[2],[1]]

出力の内側の配列1つひとつが、葉を取り除く各ラウンドに対応しています。なお、DFSの走査順序によって、同じグループ内の値の並び順は前後することがあります（例：[[4,5,3],...] と [[3,5,4],...]）。どちらの順序であっても、グルーピングとしては正しい結果です。

計算量について

この手法では各ノードを一度だけ訪問すればよいため、時間計算量は O(N)（N はノード数）、必要な記憶領域も O(N) に抑えられます。毎回木全体を走査しながら葉を削除していく素朴な方法では O(N²) かかる可能性がありますが、高さを事前に計算しておくこのアプローチなら線形時間で処理できる点が大きな利点です。