LeetCode 230. Kth Smallest Element in a BST

Problem Description

• 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?
• address

problem analysis

• This question is the same as the k-th node of the offer-binary search tree, find the k-th smallest node of BST.
• Using the result of BST in-order traversal is the property of ascending sequence, and it is easy to find the Kth .
  
  • non-recursive method
  • recursion
• The Follow up of this question is more interesting, it is how the binary tree often performs insertion and deletion operations, and often performs the operation of finding the kth element, so how to quickly get to the kth smallest element?
  Change the structure of the node, each node stores the number of all nodes in its left subtree at the same time, so that according to the size of k and the current root node, it is possible to choose whether it is the current root or the kth element in the left subtree. Recursion of the left subtree, or the kth element is in the right subtree, requires recursion of the right subtree. log(N) time complexity.

Code

• non-recursive method

public int kthSmallest(TreeNode root, int k) {
        if (root == null || k <= 0) {
            return 0;
        }
        Stack<TreeNode> stack = new Stack<>();
        TreeNode curNode = root;
        while (curNode != null || ! stack.isEmpty()) {
            while (curNode != null) {
                stack.push(curNode);
                curNode = curNode.left;
            }
            curNode = stack.pop();
            //看是否到第k个节点了
            if (--k == 0) {
                return curNode.val;
            }
            curNode = curNode.right;
        }
        return 0;
    }

• recursion

 public int kthSmallest(TreeNode root, int k) {
        if (root == null || k < 0) {
            return 0;
        }
        int[] count = new int[1]; 
        count[0] = 0;
        TreeNode res = kthSmallest(root, count, k);
        //若找到第k个节点，返回该值
        return res == null ? 0 : res.val;
    }

    // 在已经查找到了count[0]个节点之后，查找第K个节点,若找到了返回该节点，找不到，返回null
    public TreeNode kthSmallest(TreeNode root, int[] count, int k) {
        if (root == null) {
            return null;
        }
        //在左子树查找
        TreeNode leftRes = kthSmallest(root.left, count, k);
        if (leftRes != null) {//已经找到第k个节点
            return leftRes;
        }
        //尚不足k个
        if (++count[0] == k) {//当前节点
            return root;
        }
        //右子树查找
        return kthSmallest(root.right, count, k);
    }