使用Morris中序遍历,时间复杂度O(n),空间O(1)。
因为被交换的两个元素不一定相邻,所以可能需要多次swap才能恢复。此处应用的策略是记录第一次Invalid的第一个元素和最后一次Invalid的第二个元素,然后最后做一次swap。
class Solution {
public:
void recoverTree(TreeNode *root) {
TreeNode *cur = root, *prev = NULL;
pair<TreeNode*, TreeNode*> broken;
while(cur)
{
if(cur->left == NULL)
{
//visit cur
detect(broken, prev, cur);
prev = cur;
cur = cur->right;
}else
{
TreeNode *tmp = cur->left;
while(tmp->right != NULL && tmp->right != cur)
tmp = tmp->right;
if(tmp->right == NULL)
{
tmp->right = cur;
cur = cur->left;
}
else
{
//visit cur
detect(broken, prev, cur);
tmp->right = NULL;
prev = cur;
cur = cur->right;
}
}
}
swap(broken.first->val, broken.second->val);
}
void detect(pair<TreeNode*, TreeNode*> &broken, TreeNode *prev, TreeNode *cur)
{
if(prev != NULL && prev->val > cur->val)
{
//record the first element of the first invalid and the second element of the last invalid
if(broken.first == NULL) broken.first = prev;
broken.second = cur;
}
}
};