[LeetCode] Problem 230 - Kth Smallest Element in a BST

Posted on | In ,

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.

Example

No.1

Input: root = [3,1,4,null,2], k = 1

1
2
3
4
5
  3
 / \
1   4
 \
  2

Output: 1

No.2

Input: root = [5,3,6,2,4,null,null,1], k = 3

1
2
3
4
5
6
7
      5
     / \
    3   6
   / \
  2   4
 /
1

Output: 3

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?

Code

1
2
3
4
5
6
public class TreeNode {
    int val;
    TreeNode left;
    TreeNode right;
    TreeNode(int x) { val = x; }
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
private int count = 0;
private int num = 0;

public int kthSmallest(TreeNode root, int k) {
    inorder(root, k);
    return num;
}

private void inorder(TreeNode root, int k) {
    if (root.left != null)
        inorder(root.left, k);

    if (++count == k) {
        num = root.val;
        return;
    }

    if (root.right != null)
        inorder(root.right, k);
}