230. Kth Smallest Element in a BST (Medium)

Given a binary search tree, write a functionkthSmallestto find the kth smallest element in it.

Note:
You may assume k is always valid, 1 ≤ k ≤ BST's total elements.

Follow up:
What if the BST is modified (insert/delete operations) often and you need to find the kth smallest frequently? How would you optimize the kthSmallest routine?

Solution 1: Iterative, DFS O(k); O(h)

中序遍历，即可以得到递增序列，从而可以找到第k大的元素。时间复杂度O(k)。

    /**
     * Iterative, DFS O(k);O(h)
     * BT Inorder Traversal Template
     */
    public int kthSmallestB(TreeNode root, int k) {
        Stack<TreeNode> stack = new Stack<>();
        while (!stack.empty() || root != null) {
            while (root != null) {
                stack.push(root);
                root = root.left;
            }
            root = stack.pop();
            if (--k == 0) {
                break;
            }
            root = root.right;
        }
        return root.val;
    }

Solution 2: Recursive, DFS O(k); O(h)

    public int kthSmallest(TreeNode root, int k) {
        int count = countNodes(root.left);
        if (k <= count) {
            return kthSmallest(root.left, k);
        } else if (k > count + 1) {
            return kthSmallest(root.right, k - count - 1);
        }
        return root.val;
    }

    private int countNodes(TreeNode root) {
        if (root == null) {
            return 0;
        }
        return 1 + countNodes(root.left) + countNodes(root.right);
    }

Follow Up:

1.优化成constant space complexity，面试官不算递归栈空间，用递归就行

2.原题的follow up
这个问题其实是个历史问题，原题有修改，修改前提示：可以修改BST节点，能修改节点就好办了，给节点加个属性（左子树节点数），然后就好算了，复杂度降为O(h)

如果能够修改节点的数据结构TreeNode，可以增加一个字段leftCnt，表示左子树的节点个数。设当前节点为root，若 k == root.leftCnt+1, 则返回root 若 k > node.leftCnt, 则 k -= root.leftCnt+1, root=root.right 否则，node = node.left