C++で二分木が高さバランスされているかどうかを判定するプログラム

C++で二分木の高さバランスを判定する方法

二分木が与えられたとき、その木が「高さバランス」されているかどうかを判定する必要があります。

高さバランスされた木とは、すべてのノードにおいて、左部分木の高さと右部分木の高さの絶対差が0または1以内に収まっている木のことです。

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

この場合、出力は True(バランスされている)となります。

解決のアプローチ

この問題を解くには、DFS(深さ優先探索)を利用して各ノードの部分木の高さを計算し、その差を確認します。手順は以下の通りです。

• dfs() 関数を定義し、ノードを引数として受け取ります。
• ノードがNULLの場合は 0 を返します(高さ0)。
• 左部分木の高さ l を「1 + dfs(左の子ノード)」として計算します。
• 右部分木の高さ r を「1 + dfs(右の子ノード)」として計算します。
• |l - r| > 1 となる場合、フラグ変数 ret を false に設定します。
• l と r の最大値を返します(現在のノードを根とする部分木の高さ)。
• メイン処理では、まず ret = true で初期化し、根ノードに対して dfs(root) を実行した後、ret を返します。

C++での実装例

以下の実装例を見ると、理解がより深まるでしょう。

#include <bits/stdc++.h>
using namespace std;
class TreeNode {
    public:
    int val;
    TreeNode *left;
    TreeNode *right;
    TreeNode(int x) : val(x), left(NULL), right(NULL) {}
};
class Solution {
    public:
    bool ret;
    int dfs(TreeNode* node){
        if(!node)
            return 0;
        int l = 1 + dfs(node->left);
        int r = 1 + dfs(node->right);
        if(abs(l - r) > 1)
            ret = false;
        return max(l, r);
    }
    bool isBalanced(TreeNode* root) {
        ret = true;
        dfs(root);
        return ret;
    }
};
main(){
    Solution ob;
    TreeNode *root = new TreeNode(25);
    root->left = new TreeNode(19);
    root->right = new TreeNode(4);
    root->left->left = new TreeNode(9);
    root->left->right = new TreeNode(7);
    cout << (ob.isBalanced(root));
}

入力

TreeNode *root = new TreeNode(25);
root->left = new TreeNode(19);
root->right = new TreeNode(4);
root->left->left = new TreeNode(9);
root->left->right = new TreeNode(7);

出力

1

この結果 1(true)は、与えられた二分木が高さバランスされていることを示しています。