leetcode230. 二叉搜索树中第K小的元素

题目描述

给定一个二叉搜索树的根节点 root ，和一个整数 k ，请你设计一个算法查找其中第 k 个最小元素（从 1 开始计数）。

输入：root = [3,1,4,null,2], k = 1
输出：1

递归

	private int count = 0;
    private int ans = -1;

    public void dfs(TreeNode root, int k) {
        if (root == null || ans == -1) return;
        kthSmallest(root.left, k);
        if (++count == k) {
            ans = root.val;
        }
        kthSmallest(root.right, k);
    }

    public int kthSmallest(TreeNode root, int k) {
        dfs(root, k);
        return ans;
    }

迭代

迭代需要保存前一个节点的信息，所以采用栈来保存

public int kthSmallest(TreeNode root, int k) {
        ArrayDeque<TreeNode> stack = new ArrayDeque<>();
        while(root != null || !stack.isEmpty()){
            while(root != null){
                stack.push(root);
                root = root.left;
            }
            root = stack.pop();
            if( --k == 0){
                break;
            }
            root = root.right;
        }
        return root.val;
    }

通俗解法，寻找无规律树的k个最小节点，大顶堆

public static int kthSmallest(TreeNode root, int k) {
        PriorityQueue<Integer> heap = new PriorityQueue<>((x, y) -> y - x);
        ArrayDeque<TreeNode> list = new ArrayDeque<>();
        list.add(root);
        while (!list.isEmpty()){
            TreeNode cur = list.poll();
            if(heap.size() < k){
                heap.add(cur.val);
            }else if(cur.val < heap.peek()){
                heap.poll();
                heap.add(cur.val);
            }
            if(cur.left != null) list.add(root.left);
            if(cur.right != null) list.add(root.right);
        }
        return heap.peek();
    }

进阶

如果此平衡二叉树经常添加和删除，如何更快的查找kth？
方法就是保存当前节点左子树和右子树的节点个数，利用二分查找即可

public int kthSmallest(TreeNode root, int k) {
   countTree countTree = new countTree(root);
   return countTree.kthSmallest(root,k);
}

public class countTree {
   Map<TreeNode, Integer> countMap;

   public countTree(TreeNode root) {
       this.countMap = new HashMap<>();
       initCount(root);
   }

   public int initCount(TreeNode root) {
       if (root == null) return 0;
       int count = 1 + initCount(root.left) + initCount(root.right);
       countMap.put(root, count);
       return count;
   }

   public int kthSmallest(TreeNode root, int k) {
       while (root != null) {
           int count = getCount(root.left);
           if (count == k - 1) {
               break;
           } else if (count < k - 1) {
               root = root.right;
               k -= (count + 1);
           }else{
               root = root.left;
           }
       }
       return root.val;
   }

   public int getCount(TreeNode root) {
       return countMap.getOrDefault(root, 0);
   }
}