My code is as follows. Basically I created a vector to record the height of all the leaves. If the length of this vector is smaller than 2, which means we have at most one leaf node, the tree is balanced. Otherwise, find the difference of the max height & min height. If it is less than or equal to 1, the tree is balanced. For the rest cases, the tree is imbalanced.

```
void allHeights(TreeNode *root, vector<int>& record, int cur_height) {
if (root == NULL) return;
if (root->left == NULL && root->right == NULL) {
record.push_back(cur_height);
} else {
if (root->left != NULL) allHeights(root->left, record, cur_height+1);
if (root->right != NULL) allHeights(root->right, record, cur_height+1);
}
}
bool isBalanced(TreeNode *root) {
vector<int> record;
int cur_height = 0;
allHeights(root, record, cur_height);
if (record.size() <= 1) return true;
int max_height = record[0];
int min_height = record[0];
for (int i=1; i<record.size(); ++i) {
if (record[i] > max_height) max_height = record[i];
if (record[i] < min_height) min_height = record[i];
}
if (max_height - min_height > 1) return false;
else return true;
}
```