[LeetCode] 230. 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?

Thought process:
Use instance variables to keep track of k and result. Do an in-order traversal of the BST. As I traverse it in ascending order, decrement k for each node.

Solution 1 (In-order traversal):
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */
class Solution {
    private int number;
    private int count;
    
    public int kthSmallest(TreeNode root, int k) {
        count = k;
        inorder(root);
        return number;
    }
    
    private void inorder(TreeNode node) {
        if (node.left != null) {
            inorder(node.left);
        }
        count--;
        if (count == 0) {
            number = node.val;
            return;
        }
        if (node.right != null) {
            inorder(node.right);
        }
    }
}

This can also be solved using binary search. Count the number of nodes in root's left and right sub-trees. 
  1. If k <= left: recurse on the left sub-tree.
  2. if k > left + 1: recurse on the right sub-tree.
  3. Otherwise, root is the kth smallest node.

Solution 2 (Binary search):
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
/**
 * 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) {
        int left = countNodes(root.left);
        if (k <= left) {
            return kthSmallest(root.left, k);
        }
        if (k > left + 1) {
            return kthSmallest(root.right, k - left - 1);
        }
        return root.val;
    }
    
    private int countNodes(TreeNode node) {
        if (node == null) {
            return 0;
        }
        return countNodes(node.left) + 1 + countNodes(node.right);
    }
}

If the BST is modified often, then I could add fields into TreeNode to keep track of number of nodes in the left and right sub-trees for each node. As the tree is modified, also update the counts. Then I can use binary search to get the kth smallest number quickly.

Time complexity: worst case O(k) for a linear tree.

Comments

Popular posts from this blog

[LeetCode] 714. Best Time to Buy and Sell Stock with Transaction Fee

[LeetCode] 269. Alien Dictionary

[LeetCode] 631. Design Excel Sum Formula