《Cracking the Coding Interview》——第4章：树和图——题目1

2014-03-19 03:30

题目：判断一个二叉树是否为平衡二叉树，即左右子树高度相差不超过1。

解法：递归算高度并判断即可。

代码：

// 4.1 Implement an algorithm to check if a bianry tree is height-balanced.
#include <algorithm>
#include <cstdio>
#include <unordered_map>
using namespace std;

struct TreeNode {
    int val;
    TreeNode *left;
    TreeNode *right;
    
    TreeNode(int _val = 0): val(_val), left(nullptr), right(nullptr) {};
};

void constructBinaryTree(TreeNode *&root)
{
    int val;
    
    scanf("%d", &val);
    if (val <= 0) {
        root = nullptr;
    } else {
        root = new TreeNode(val);

        constructBinaryTree(root->left);
        constructBinaryTree(root->right);
    }
}

void clearBinaryTree(TreeNode *&root) {
    if (root == nullptr) {
        return;
    } else {
        clearBinaryTree(root->left);
        clearBinaryTree(root->right);
        delete root;
        root = nullptr;
    }
}

void calcHeights(TreeNode *root, unordered_map<TreeNode *, int> &heights)
{
    if (root == nullptr) {
        heights[root] = 0;
    } else {
        calcHeights(root->left, heights);
        calcHeights(root->right, heights);
        heights[root] = max(heights[root->left], heights[root->right]) + 1;
    }
}

bool isBalanced(TreeNode *root, unordered_map<TreeNode *, int> &heights)
{
    if (root == nullptr) {
        return true;
    } else {
        return abs(heights[root->left] - heights[root->right]) <= 1;
    }
}

int main()
{
    TreeNode *root;
    unordered_map<TreeNode *, int> heights;
    
    while (true) {
        constructBinaryTree(root);
        if (root == nullptr) {
            break;
        }
        
        calcHeights(root, heights);
        if (isBalanced(root, heights)) {
            printf("Yes
");
        } else {
            printf("No
");
        }
        heights.clear();
        clearBinaryTree(root);
    }
    
    return 0;
}