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.

Example 1:

Input: root = [3,1,4,null,2], k = 1
   3
  / \
 1   4
  \
   2
Output: 1

Example 2:

Input: root = [5,3,6,2,4,null,null,1], k = 3
       5
      / \
     3   6
    / \
   2   4
  /
 1
Output: 3

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?

time o(n) o(h+k), space o(h)

inorder traversal find kth

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     int left;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 *     
 * }
 */
class Solution {
    public int kthSmallest(TreeNode root, int k) {
        //first thought was inorder traversal
        if(root == null)
            return -1;
        
        Stack<TreeNode> stack = new Stack<>();
        
        TreeNode cur = root;
        
        
        while(cur != null || !stack.isEmpty()){
            while(cur != null){
                stack.push(cur);
                cur = cur.left;
            }
            
            cur = stack.pop();
            
            k--;
                
            if(k == 0)
                return cur.val;
            
            cur = cur.right;
        }
        
        return -1;
    }
}

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?

再用inorder traversal 效率会低

If I am allowed to add a new member to the node, I would probably add a pointer to each node so that it always knows its parent. With this pointer, it could be much easier to find the next element. I will always hold a pointer P to the kth element. If things change in elements greater than the kth element, take no action. If an element in the first kth element is deleted, we find the next element of P. If a new element less than P is added, find the previous element of P.

parent pointer

quick select

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */
class Solution {
    public int kthSmallest(TreeNode root, int k) {
            if(root == null){
                return -1;
            }
            Map<TreeNode, Integer> map = new HashMap<>();
        
        
            countNode(root,map);
        
            return quickSelect(root, map, k);
    }
    
    public int countNode(TreeNode root, Map<TreeNode, Integer> map ){
        if(root == null)
            return 0;
        
        int left = countNode(root.left, map);
        int right = countNode(root.right, map);
        
        map.put(root, left+right+1);
        
        return left+right+1;
    }
    
    public int quickSelect(TreeNode root, Map<TreeNode,Integer> map, int k){
        if(root == null)
            return -1;
        
        int left = root.left == null ? 0 : map.get(root.left);
        
        if(left >= k){
            return quickSelect(root.left, map,k);
        }
        
        if(left +1 == k){
            return root.val;
        }
        
        return quickSelect(root.right, map, k - left - 1);
    }
}