Kth Smallest Element in a BST

Given a binary search tree, write a function kthSmallest to 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?

Hint:

Try to utilize the property of a BST. What if you could modify the BST node’s structure? The optimal runtime complexity is O(height of BST).

中序遍历

复杂度

时间 O(N) 空间 O(N)

思路

因为左节点小于根节点小于右节点，二叉搜索树的一个特性就是中序遍历的结果就是树内节点从小到大顺序输出的结果。这里采用迭代形式，我们就可以在找到第k小节点时马上退出。中序遍历详见Binary Tree Traversal。

代码

public class Solution {
    public int kthSmallest(TreeNode root, int k) {
        Stack<TreeNode> s = new Stack<TreeNode>();
        while(root!=null){
            s.push(root);
            root = root.left;
        }
        while(!s.isEmpty()){
            TreeNode curr = s.pop();
            k--;
            if(k == 0) return curr.val;
            if(curr.right != null){
                curr = curr.right;
                while(curr!=null){
                    s.push(curr);
                    curr = curr.left;
                }
            }
        }
        return 0;
    }
}

后续 Follow Up

这题的难点其实在于Follow Up：如果我们频繁的操作该树，并且频繁的调用kth函数，有什么优化方法使时间复杂度降低至O(h)？h是树的高度。根据提示，我们可以在TreeNode中加入一个rank成员，这个变量记录的是该节点的左子树中节点的个数，其实就是有多少个节点比该节点小。这样我们就可以用二叉树搜索的方法来解决这个问题了。这个添加rank的操作可以在建树的时候一起完成。