C++で二分木の「良いノード」を数えるアルゴリズム

問題の概要

二分木が与えられたとき、木の中のノードXは、ルートからXへの経路上に、Xより大きい値を持つノードが存在しない場合に「良いノード(good node)」と定義されます。この記事では、二分木に含まれる良いノードの数を求める方法を解説します。

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

この場合、出力は 4 となります。色付きで示されたノードが「良いノード」に該当します。

解決アプローチ

この問題は、深さ優先探索(DFS)を使うことで効率的に解くことができます。具体的な手順は以下の通りです。

DFS関数の定義

• dfs() 関数を定義します。引数としてノード(node)と現在までの最大値(val)を受け取ります。
• node が null の場合は何もせず return します。
• val が現在のノードの値以下であれば、答え(ret)に1を加算します。そうでなければ加算しません。
• dfs() を再帰的に呼び出し、左の子ノードには val とノードの値の大きい方を渡します。
• 同様に、右の子ノードに対しても再帰的に dfs() を呼び出します。

メイン処理の流れ

• カウンタ変数 ret を 0 で初期化します。
• dfs(root, -∞) を実行します。初期値に負の無限大を設定することで、ルートノードは必ず「良いノード」と判定されます。
• 最後に ret を返します。

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 ret;
   void dfs(TreeNode* node, int val){
      if (!node)
         return;
      ret += val <= node->val;
      dfs(node->left, max(val, node->val));
      dfs(node->right, max(val, node->val));
   }
   int goodNodes(TreeNode* root){
      ret = 0;
      dfs(root, INT_MIN);
      return ret;
   }
};
main(){
   Solution ob;
   vector<int> v = {3,1,4,3,NULL,1,5};
   TreeNode *root = make_tree(v);
   cout << (ob.goodNodes(root));
}

入力例

{3,1,4,3,NULL,1,5}

出力

4

まとめ

このアルゴリズムは、各ノードを一度だけ訪問するため、計算量はO(N)(Nはノード数)、再帰によるスタック使用を含めた空間計算量は最悪情况下O(N)となります。DFSで経路上の最大値を引き継ぐというシンプルな発想により、効率的に「良いノード」を判定できる点がポイントです。