[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.
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?
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.
- If k <= left: recurse on the left sub-tree.
- if k > left + 1: recurse on the right sub-tree.
- 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
Post a Comment