Find kth smallest element in BST 找第k个小(大)的值在BST中

今天geeksforgeeks上刷出来的题.一看不能用stack, 基本就是Morris遍历了. 但是做了一下, 发现需要注意的地方挺多的, 比如当我们算visit的时候, 在rebuild link后, 我们也需要increase counter. 不然打印出来的是The Bottom of Tree

代码:

 public static int find(TreeNode root, int k) {
        int count = 0;//
        int res = Integer.MAX_VALUE;
        TreeNode cur = root;
        while (cur != null) {
            if (cur.left == null) { // if left == null, then we print(or do something, in here we increase counter)
                count++;
                if (count == k) // if it is equal to k
                    res = cur.val; // set the res
                cur = cur.right; // move right
            }
            else { // if left != null,in other words, we need to build the link
                TreeNode pre = cur.left; //make a pointer move left
                while (pre.right != null && pre.right != cur) // use the pointer to find the predecessor of cur
                    pre = pre.right; // move right
                if (pre.right == null) { // if the right is null, build link,
                    pre.right = cur;// set right to cur
                    cur = cur.left;// move left
                }
                else { //if right != null, restore the link
                    pre.right = null;//set right to nul
                    count++; // because we rebuild our old bst, we need to increase counter
                    if (count == k) //if we do not increase counter in here, the result will be THE BOTTOM OF TREE. because after rebuild the link, the node is already visited.
                        res = cur.val;
                    cur = cur.right; // move left.
                }
            }
        }
        return res;
    }