You are given the root node of a binary search tree and an integer K. Your task is to find the Kth largest/smallest element in Binary Search Tree. You need to return an array containing two integers, where the first integer is the Kth smallest value and the second integer is the Kth largest value from all the nodes present in the tree.
Example 1
Input: Root of tree with nodes [5, 3, 6, 2, 4, null, null, 1], K = 3
Output: [3, 4]
Explanation: The sorted values of the tree are [1, 2, 3, 4, 5, 6]. The 3rd smallest value is 3 from the beginning, and the 3rd largest value is 4 from the end.
Example 2
Input: Root of tree with nodes [2, 1, 3], K = 2
Output: [2, 2]
Explanation: The sorted values are [1, 2, 3]. The 2nd smallest value is 2, and the 2nd largest value is also 2.
Brute Force Approach
Think of organizing a messy stack of test papers based on student scores. If we lay out all the papers in increasing order of their marks, finding the Kth lowest or Kth highest score simply requires picking the Kthpaper from the start or the end of the line.
A binary search tree has a special property where traversing it in an Inorder manner (visiting the left child, then the root, then the right child) yields all the node values in strictly increasing order. By doing this traversal and writing down every value in an array, we get a fully sorted list of elements. From this sorted list, the Kthsmallest element is found at index K minus 1. The Kth largest element is found by counting K positions backward from the end of the list.
Algorithm
Create an empty array to store the node values.
Start an inorder traversal from the root node to get the node values in sorted order.
Recursively visit the left subtree first so smaller values gets added first.
Append the current node's value to the array, all smaller values than root added now root.
Recursively visit the right subtree, now larger values than root are added.
Once the traversal is complete and the array is filled, find the Kth smallest element at index K minus 1.
Find the Kth largest element at the index equal to the total size of the array minus K.
Return these two values in a new array.
Dry Run
Kth Smallest and Largest Element in BST Brute Dry Run
Solution
#include <bits/stdc++.h>using namespace std;// Definition for a binary tree nodestruct TreeNode { int val; TreeNode *left; TreeNode *right; TreeNode(int x) : val(x), left(NULL), right(NULL) {}};class Solution {private: // Helper function to perform inorder traversal and store values void inorder(TreeNode* root, vector<int>& values) { // Base condition to stop recursion when a leaf's child is reached if (root == NULL) { return; } // Recursively traverse the left subtree inorder(root->left, values); // Store the value of the current node values.push_back(root->val); // Recursively traverse the right subtree inorder(root->right, values); }public: // Main function to find the Kth smallest and largest elements vector<int> findKthElements(TreeNode* root, int k) { vector<int> values; // Populate the values array using inorder traversal inorder(root, values); int n = values.size(); vector<int> result(2); // Retrieve the Kth smallest element from the beginning result[0] = values[k - 1]; // Retrieve the Kth largest element from the end result[1] = values[n - k]; return result; }};// Driver code starts hereint main() { // Constructing a sample binary search tree TreeNode* root = new TreeNode(2); root->left = new TreeNode(1); root->right = new TreeNode(3); Solution obj; int k = 2; // Calling the main function to get the result vector<int> result = obj.findKthElements(root, k); // Printing the output cout << "Smallest: " << result[0] << ", Largest: " << result[1] << endl; return 0;}Complexity Analysis
Time Complexity: O(N), where N is the number of nodes. We visit every node exactly once during the inorder traversal.
Space Complexity: O(N), because we store all the node values in an external array before fetching the results.
Optimal Approach
The previous approach requires storing every single element in a list, which takes up a lot of extra space. If a tree has millions of nodes, storing all of them just to find the second or third element is inefficient. We can solve this by skipping the storage entirely.
Instead of making all students stand in a line and taking up a massive hallway, we just stand at the door. As students walk out in increasing order of their marks, we keep a tally on a notepad. When the Kth student walks out, we write down their name and we can stop counting.
For the binary search tree, we can traverse it in the standard inorder way (Left, Root, Right) to count the smallest elements. When our counter reaches K, we have found the Kth smallest element. To find the Kth largest element, we can do the exact same thing but in reverse. By traversing Right, Root, Left, we process the elements in decreasing order. When the counter reaches K during this reverse traversal, we have found the Kth largest element.
If you want to understand this in detail, check our guide on Tree Traversals.
Algorithm
Create two variables to keep track of our count, one for the smallest and one for the largest.
Create two variables to store the final answers.
Write a recursive function for the Kth smallest element that visits the left node, increments the counter, checks if the counter equals K, stores the value if true, and then visits the right node.
Write a second recursive function for the Kth largest element that visits the right node first, increments the counter, checks if the counter equals K, stores the value if true, and then visits the left node.
Call both functions sequentially.
Return the stored answers in an array.
Dry Run
Kth Smallest and Largest Element in BST
Solution
#include <bits/stdc++.h>using namespace std;// Definition for a binary tree nodestruct TreeNode { int val; TreeNode *left; TreeNode *right; TreeNode(int x) : val(x), left(NULL), right(NULL) {}};class Solution {private: int countSmallest = 0; int countLargest = 0; int kthSmallestVal = -1; int kthLargestVal = -1; // Standard inorder traversal to find the smallest element void findSmallest(TreeNode* root, int k) { if (root == NULL || countSmallest >= k) { return; } // Process left subtree first findSmallest(root->left, k); // Increment count and check if current node is the target countSmallest++; if (countSmallest == k) { kthSmallestVal = root->val; return; } // Process right subtree findSmallest(root->right, k); } // Reverse inorder traversal to find the largest element void findLargest(TreeNode* root, int k) { if (root == NULL || countLargest >= k) { return; } // Process right subtree first to get descending order findLargest(root->right, k); // Increment count and check if current node is the target countLargest++; if (countLargest == k) { kthLargestVal = root->val; return; } // Process left subtree findLargest(root->left, k); }public: // Main function to coordinate both traversals vector<int> findKthElements(TreeNode* root, int k) { findSmallest(root, k); findLargest(root, k); vector<int> result(2); result[0] = kthSmallestVal; result[1] = kthLargestVal; return result; }};// Driver code starts hereint main() { TreeNode* root = new TreeNode(5); root->left = new TreeNode(3); root->right = new TreeNode(6); root->left->left = new TreeNode(2); root->left->right = new TreeNode(4); Solution obj; int k = 3; vector<int> result = obj.findKthElements(root, k); cout << "Smallest: " << result[0] << ", Largest: " << result[1] << endl; return 0;}Complexity Analysis
Time Complexity: O(N), because in the worst-case scenario where K is equal to the total number of nodes, we might traverse all nodes. However, on average, we stop early once the counter reaches K.
Space Complexity: O(H), where H is the height of the tree. This space is strictly due to the recursion stack during the depth-first traversal. No external arrays are used to store all elements.
Interview follow-up Questions
Usually, problems specify that K is always valid based on the tree size. If it is not, our optimized approach will safely return the default initialized value since the counter will never reach K.
Be the first to add a comment.