Problem Description:
Given a non-empty binary search tree and a target value, find k values in the BST that are closest to the target.
Note:
- Given target value is a floating point.
- You may assume k is always valid, that is: k ≤ total nodes.
- You are guaranteed to have only one unique set of k values in the BST that are closest to the target.
Follow up:
Assume that the BST is balanced, could you solve it in less than O(n) runtime (where n = total nodes)?
Hint:
- Consider implement these two helper functions:
-
getPredecessor(N), which returns the next smaller node to N. -
getSuccessor(N), which returns the next larger node to N.
-
- Try to assume that each node has a parent pointer, it makes the problem much easier.
- Without parent pointer we just need to keep track of the path from the root to the current node using a stack.
- You would need two stacks to track the path in finding predecessor and successor node separately.
There is a very simple idea to keep a max heap of size k and elements in the heap are sorted by their absolute difference from target. For the max heap, we will simply use the default priority_queue. Then we simply traverse the tree and push all its nodes to the heap while maintaining the size of heap not larger than k.
The code is as follows.
1 class Solution { 2 public: 3 vector<int> closestKValues(TreeNode* root, double target, int k) { 4 priority_queue<pair<double, int>> pq; 5 closestK(root, pq, target, k); 6 vector<int> closest; 7 while (!pq.empty()) 8 closest.push_back(pq.top().second), pq.pop(); 9 return closest; 10 } 11 private: 12 void closestK(TreeNode* node, priority_queue<pair<double, int>>& pq, double target, int k) { 13 if (!node) return; 14 pq.push(make_pair(abs(target - node -> val), node -> val)); 15 if (pq.size() > k) pq.pop(); 16 closestK(node -> left, pq, target, k); 17 closestK(node -> right, pq, target, k); 18 } 19 };
Well, the hints of the problem suggest a solution using two stacks to track the predecessors and successors. This post has shared a nice implementation of it. The code is rewritten in C++ as follows.
1 class Solution { 2 public: 3 vector<int> closestKValues(TreeNode* root, double target, int k) { 4 vector<int> closest(k); 5 stack<int> pre, suc; 6 inorder(root, target, false, pre); 7 inorder(root, target, true, suc); 8 for (int i = 0; i < k; i++) { 9 if (pre.empty()) closest[i] = suc.top(), suc.pop(); 10 else if (suc.empty()) closest[i] = pre.top(), pre.pop(); 11 else if (abs(target - pre.top()) < abs(target - suc.top())) 12 closest[i] = pre.top(), pre.pop(); 13 else closest[i] = suc.top(), suc.pop(); 14 } 15 return closest; 16 } 17 private: 18 void inorder(TreeNode* root, double target, bool reversed, stack<int>& s) { 19 if (!root) return; 20 inorder(reversed ? root -> right : root -> left, target, reversed, s); 21 if ((reversed && root -> val <= target) || (!reversed && root -> val > target)) return; 22 s.push(root -> val); 23 inorder(reversed ? root -> left : root -> right, target, reversed, s); 24 } 25 };
