Question
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).
Solution 1 -- Inorder Traversal
Again, we use the feature of inorder traversal of BST. But this solution is not best for follow up. Time complexity O(n), n is the number of nodes.
/**
* 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) {
TreeNode current = root;
Stack<TreeNode> stack = new Stack<TreeNode>();
while (current != null || !stack.empty()) {
if (current != null) {
stack.push(current);
current = current.left;
} else {
TreeNode tmp = stack.pop();
k--;
if (k == 0) {
return tmp.val;
}
current = tmp.right;
}
}
return -1;
}
}
Solution 2 -- Augmented Tree
The idea is to maintain rank of each node. We can keep track of elements in a subtree of any node while building the tree. Since we need K-th smallest element, we can maintain number of elements of left subtree in every node.
Assume that the root is having N nodes in its left subtree. If K = N + 1, root is K-th node. If K < N, we will continue our search (recursion) for the Kth smallest element in the left subtree of root. If K > N + 1, we continue our search in the right subtree for the (K – N – 1)-th smallest element. Note that we need the count of elements in left subtree only.
Time complexity: O(h) where h is height of tree.
(referrence: GeeksforGeeks)
Here, we construct tree in a way that is taught during Algorithm class.
"size" is an attribute which indicates number of nodes in sub-tree rooted in that node.
Time complexity: constructing tree O(n), find Kth smallest number O(h).
start:
if K = root.leftElement + 1
root node is the K th node.
goto stop
else if K > root.leftElements
K = K - (root.leftElements + 1)
root = root.right
goto start
else
root = root.left
goto srart stop
/**
* Definition for a binary tree node.
* public class TreeNode {
* int val;
* TreeNode left;
* TreeNode right;
* TreeNode(int x) { val = x; }
* }
*/
class ImprovedTreeNode {
int val;
int size; // number of nodes in the subtree that rooted in this node
ImprovedTreeNode left;
ImprovedTreeNode right;
public ImprovedTreeNode(int value) {val = value;}
} public class Solution { // Construct ImprovedTree recursively
public ImprovedTreeNode createAugmentedBST(TreeNode root) {
if (root == null)
return null;
ImprovedTreeNode newHead = new ImprovedTreeNode(root.val);
ImprovedTreeNode left = createAugmentedBST(root.left);
ImprovedTreeNode right = createAugmentedBST(root.right);
newHead.size = 1;
if (left != null)
newHead.size += left.size;
if (right != null)
newHead.size += right.size;
newHead.left = left;
newHead.right = right;
return newHead;
} public int findKthSmallest(ImprovedTreeNode root, int k) {
if (root == null)
return -1;
ImprovedTreeNode tmp = root;
int leftSize = 0;
if (tmp.left != null)
leftSize = tmp.left.size;
if (leftSize + 1 == k)
return root.val;
else if (leftSize + 1 > k)
return findKthSmallest(root.left, k);
else
return findKthSmallest(root.right, k - leftSize - 1);
} public int kthSmallest(TreeNode root, int k) {
if (root == null)
return -1;
ImprovedTreeNode newRoot = createAugmentedBST(root);
return findKthSmallest(newRoot, k);
}
}