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).

Tips:

The first solution is actually the worst, but good for the follow up. it is an O(n^2) worst, O(nlogn) average case algorithm.

the second one is O(n) but it doesn't terminate when solution is found, instead it continues to scan the remaining nodes which is a waste without question.

The third one is a traditional iterative in order traversal and should solve the problem.

Follow Up:

Hi I think the first solution is for the follow up question. To find the kth element in the BST, we'll definitely need o(k), that is o(n) complexity. But if we are allowed to modify the node structure of BST, we can simply store the count value in every node. And thus we can achieve the search in o(logn)

Code:

解法1：Binary Search (dfs): most preferable

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */
public class Solution {
    public int kthSmallest(TreeNode root, int k) {
        int count = countNodes(root.left);
        if (count == k - 1){
            return root.val;
        } else if (count + 1 < k) {
            return kthSmallest(root.right, k - count - 1);
        } else {
            return kthSmallest(root.left, k);
        }
    }
    private int countNodes(TreeNode node) {
        if (node == null) {
            return 0;
        }
        return countNodes(node.left) + countNodes(node.right) + 1;
    }
}

解法2：DFS in-order recursive:

// better keep these two variables in a wrapper class
private static int number = 0;
private static int count = 0;

public int kthSmallest(TreeNode root, int k) {
    count = k;
    helper(root);
    return number;
}

public void helper(TreeNode n) {
    if (n.left != null) helper(n.left);
    count--;
    if (count == 0) {
        number = n.val;
        return;
    }
    if (n.right != null) helper(n.right);
}

解法3：DFS in-order iterative:

public int kthSmallest(TreeNode root, int k) {
        Stack<TreeNode> st = new Stack<>();

        while (root != null) {
            st.push(root);
            root = root.left;
        }

        while (k != 0) {
            TreeNode n = st.pop();
            k--;
            if (k == 0) return n.val;
            TreeNode right = n.right;
            while (right != null) {
                st.push(right);
                right = right.left;
            }
        }

        return -1; // never hit if k is valid
  }