Leetcode 110 Solution

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
class Solution {
    class State
    {
    public:
        int stage;
        int left_height;
        int right_height;
        bool left_balanced;
        bool right_balanced;

        State(int _stage, int _left_height, int _right_height, bool _left_balanced, bool _right_balanced)
        {
            stage = _stage;
            left_height = _left_height;
            right_height = _right_height;
            left_balanced = _left_balanced;
            right_balanced = _right_balanced;
        }
    };

public:
    bool isBalanced(TreeNode* root) {
        stack<pair<TreeNode*, State>> s;

        if (root == NULL)
            return true;

        s.push(make_pair(root, State(0, 0, 0, true, true)));

        int last_height = 0;
        bool last_balanced = true;

        while (s.empty() == false)
        {
            auto pair = s.top();
            auto node = pair.first;
            auto state = pair.second;

            s.pop();

            if (node == NULL)
            {
                last_height = 0;
                last_balanced = true;
                continue;
            }

            if (state.stage == 0)
            {
                state.stage = 1;
                s.push(make_pair(node, state));

                s.push(make_pair(node->left, State(0, 0, 0, true, true)));
            }
            else if (state.stage == 1)
            {
                state.stage = 2;
                state.left_height = last_height;
                state.left_balanced = last_balanced;
                s.push(make_pair(node, state));

                s.push(make_pair(node->right, State(0, 0, 0, true, true)));
            }
            else if (state.stage == 2)
            {
                state.right_height = last_height;
                state.right_balanced = last_balanced;

                last_balanced = state.left_balanced && state.right_balanced && (state.left_height - state.right_height <= 1) && (state.left_height - state.right_height >= -1);
                last_height = max(state.left_height, state.right_height) + 1;
            }
        }

        return last_balanced;
    }
};