思路分析
思想承接上一道题目“Validate Binary Search Tree”,不过写起来真心是复杂许多,核心还是在于利用树的递归性质来求解。存在四种基本情况:交换完全存在于左/右子树中、根与左子树中的元素交换、根与右子树中的元素交换、左右子树中的元素相互交换。对于第一种情况,直接对左右子树继续DFS即可;对于第二、三种情况,交换根与对应元素;对于第四种情况,交换两个子树中的对应元素。
问题在于,如何判断当前是属于哪一种情况呢?显然还是要从二叉树的性质来入手。如果是交换完全存在于左右子树中,那么显然根节点元素大于左子树中的最大值,且小于右子树中的最小值;如果是根节点与左子树中元素发生了交换,那么就应该有左子树中的最大元素大于根节点元素,且小于右子树中最小元素;如果是左右子树中元素发生了交换,则应该有根节点元素大于右子树中的最小值,但是小于左子树中的最大值。
所以我们看到,还需要写两个递归求树中最值并返回指向对应节点指针的辅助函数。最后,把上述这些情况分类讨论一遍,再结合上二叉树的五种基本形态做一下分类以避免在边界条件上挂掉,即可AC这道蛋疼的题目。当然实际代码可以精简许多,不然也容易出错,我只是懒得再改了。
代码
class Solution {
public:
pair<int, TreeNode *> get_max(TreeNode *root) {
if ( root == NULL ) return make_pair(0, (TreeNode *)NULL);
int ans = root->val;
TreeNode * p_ans = root;
pair<int, TreeNode *> res;
if ( root->left != NULL ) {
res = get_max(root->left);
if ( res.first > ans ) {
ans = res.first;
p_ans = res.second;
}
}
if ( root->right != NULL ) {
res = get_max(root->right);
if ( res.first > ans ) {
ans = res.first;
p_ans = res.second;
}
}
return make_pair(ans, p_ans);
}
pair<int, TreeNode *> get_min(TreeNode *root) {
if ( root == NULL ) return make_pair(0, (TreeNode *)NULL);
int ans = root->val;
TreeNode * p_ans = root;
pair<int, TreeNode *> res;
if ( root->left != NULL ) {
res = get_min(root->left);
if ( res.first < ans ) {
ans = res.first;
p_ans = res.second;
}
}
if ( root->right != NULL ) {
res = get_min(root->right);
if ( res.first < ans ) {
ans = res.first;
p_ans = res.second;
}
}
return make_pair(ans, p_ans);
}
void recoverTree(TreeNode *root) {
if ( NULL == root ) return;
pair<int, TreeNode *> max_l = get_max(root->left), min_r = get_min(root->right);
if ( root->left != NULL && root->right != NULL ) {
if ( max_l.first < root->val && root->val < min_r.first ) {
//如果正常,说明错误存完全在于左/右子树中
recoverTree(root->left);
recoverTree(root->right);
} else if ( max_l.first > root->val && max_l.first < min_r.first ) {
//是否是根与左子树交换
swap(root->val, max_l.second->val);
} else if ( min_r.first < root->val && min_r.first > max_l.first ) {
//是否是根与右子树交换
swap(root->val, min_r.second->val);
} else if ( min_r.first < root->val && root->val < max_l.first ) {
//是否是左右子树之间交换
swap(max_l.second->val, min_r.second->val);
}
} else if ( root->left != NULL && root->right == NULL ) {
//处理只有左子树的情况
if ( max_l.first < root->val ) {
//如果正常
recoverTree(root->left);
} else if ( max_l.first > root->val ) {
swap(root->val, max_l.second->val);
}
} else if ( root->left == NULL && root->right != NULL ) {
//处理只有右子树的情况
if ( min_r.first > root->val ) {
//如果正常
recoverTree(root->right);
} else if ( min_r.first < root->val ) {
swap(root->val, min_r.second->val);
}
} else if ( root->left == NULL && root->right == NULL ) {
//处理单个节点的情况,直接返回
return;
}
}
};