【LeetCode题解】230. 二叉搜索树中第K小的元素

给定一个二叉搜索树，编写一个函数 kthSmallest 来查找其中第 k 个最小的元素。

说明：
你可以假设 k 总是有效的，1 ≤ k ≤ 二叉搜索树元素个数。

示例 1:

输入: root = [3,1,4,null,2], k = 1
   3
  / \
 1   4
  \
   2
输出: 1
示例 2:

输入: root = [5,3,6,2,4,null,null,1], k = 3
       5
      / \
     3   6
    / \
   2   4
  /
 1
输出: 3

来源：力扣（LeetCode）
链接：https://leetcode-cn.com/problems/kth-smallest-element-in-a-bst
著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。

常规解法：中序遍历找出第k大

进阶问题：
如果二叉搜索树经常被修改（插入/删除操作）并且你需要频繁地查找第 k 小的值，你将如何优化 kthSmallest 函数？

进阶问题尝试解：

class Solution {
    Map<TreeNode,Integer> m = new HashMap<>();
    public int kthSmallest(TreeNode root, int k) {
        /*
        如果二叉搜索树经常被修改（插入/删除操作）并且你需要频繁地查找第 k 小的值，你将如何优化 kthSmallest 函数？
        一个节点左子树的个数决定了比它小的节点个数
        可以在插入查找时向下传播将沿途每个节点的cnt++，删除操作时向上传播将它所有的父节点、父节点的父节点....root的cnt--
        */
        getCnt(root);
        int cnt = root.left==null? 0:m.get(root.left);
        while(cnt != k-1){
            if(cnt<k-1){//向右，抛弃左子树+当前节点，并在k中将他们的个数减去
                root = root.right;
                k -= (cnt+1);
            }else{
                root = root.left;
            }
            cnt = root.left==null? 0:m.get(root.left);
        }
        return root.val;
    }
    
    int getCnt(TreeNode root){//当前节点以及左右子树的节点个数之和
        if(root == null){
            return 0;
        }
        int cnt = 1;
        if(root.left!=null){
            cnt += getCnt(root.left);
        }
        if(root.right!=null){
            cnt += getCnt(root.right);
        }
        m.put(root, cnt);
        return cnt;
    }
}