【Lintcode】691. Recover Binary Search Tree

题目地址：

https://www.lintcode.com/problem/recover-binary-search-tree/description

给定一棵二叉搜索树，题目保证没有重复数字。其中有可能有两个节点交换过。要求将该树恢复原状并返回树根。如果没交换过则返回原树树根。

思路是中序遍历。可以用栈来做。先考虑如果一个从小到大排好序的数组$A$将其中两个数交换会产生什么效果，例如$(1,4,3,2)$，那么会有其中一个数比前面的数小，则两个被交换过的数，第一个一定是首先发现$A[i]>A[i+1]$中的那个$A[i]$，第二个一定是最后一个满足$A[i]>A[i+1]$中的那个$A[i+1]$。上面例子中就是$4$和$2$。所以中序遍历的同时将这两个node找到即可。如果找不到则返回原树树根，否则将两个节点的值交换。代码如下：

import java.util.Deque;
import java.util.LinkedList;

public class Solution {
    /**
     * @param root: the given tree
     * @return: the tree after swapping
     */
    public TreeNode bstSwappedNode(TreeNode root) {
        // write your code here
        TreeNode swap1 = null, swap2 = null;
        // last记录中序遍历到cur的时候上一个遍历的节点是哪个
        TreeNode cur = root, last = null;
        // first记录是不是第一次违反顺序
        boolean first = true;
        Deque<TreeNode> stack = new LinkedList<>();
        while (cur != null || !stack.isEmpty()) {
            while (cur != null) {
                stack.push(cur);
                cur = cur.left;
            }
            
            cur = stack.pop();
            // 如果cur不是第一个遍历到的节点，则将cur与其中序遍历前驱比较
            if (last != null) {
            	// 如果找到了第一次违反的地方，则进行赋值，并标记第一次为false
                if (cur.val < last.val && first) {
                    swap1 = last;
                    // 这里swap2也需要赋值，以防止交换的两个节点恰好相邻
                    swap2 = cur;
                    first = false;
                } else if (cur.val < last.val) {
                	// 如果找到了第k次违反的地方，则覆盖swap2的值
                    swap2 = cur;
                }
            }

            last = cur;
            cur = cur.right;
        }
        
        // 存在违反顺序的节点，则进行交换
        if (swap1 != null) {
            int tmp = swap1.val;
            swap1.val = swap2.val;
            swap2.val = tmp;
        }
        
        return root;
    }
}

class TreeNode {
    int val;
    TreeNode left, right;
    
    public TreeNode(int val) {
        this.val = val;
    }
}

时间复杂度$O(n)$，空间$O(h)$。