Leetcode 110 Solution

https://leetcode.com/problems/balanced-binary-tree

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;
}
};
``````