题目链接:https://leetcode.com/problems/closest-binary-search-tree-value-ii/
Given a non-empty binary search tree and a target value, find k values in the BST that are closest to the target.
Note:
Follow up:
Assume that the BST is balanced, could you solve it in less than O(n) runtime (where n = total nodes)?
Hint:
getPredecessor(N)
, which returns the next smaller node to N.getSuccessor(N)
, which returns the next larger node to N.思路: 利用二叉搜索树的中序遍历, 如果结果集合中元素还不到k个, 就把当前元素加到集合中去, 如果集合中的元素个数多于k了, 那么有二种情况:
1. target的值比集合中最小的值小, 因为中序遍历是有序的, 最小元素就是第0个元素
2. target的值大于结果集合中的最小值, 小于最大值
3. target的值大于集合中的最大值
当集合中的元素个数大于k的时候, 我们需要进行元素替换.
如果是第一种情况, 那么将无法进行替换, 因为第0个元素就是最靠近的值, 并且在集合中元素是有序的, 因此此时集合中就是最靠近的k个元素.
如果是第二三种情况, 那么我们比较当前值和集合中最小值, 如果当前结点值比那个值更靠近target, 那么我们就用当前元素替换最小的值
直到我们无法再找到能够替换的元素, 就可以返回结果了.
时间复杂度O(N)
代码如下:
/**
* Definition for a binary tree node.
* struct TreeNode {
* int val;
* TreeNode *left;
* TreeNode *right;
* TreeNode(int x) : val(x), left(NULL), right(NULL) {}
* };
*/
class Solution {
public:
vector closestKValues(TreeNode* root, double target, int k) {
if(!root) return result;
closestKValues(root->left, target, k);
if(result.size() < k) result.push_back(root->val);
else if(fabs(target-result[index]) > fabs(target-root->val))
{
result[index++] = root->val;
if(index==k) index = 0;
}
else return result;
closestKValues(root->right, target, k);
return result;
}
private:
vector result;
int index = 0;
};
参考: https://leetcode.com/discuss/83431/c-simple-inorder-solution-with-o-n-runtime-and-o-k-memory
还有一种O(logn)的方法, 利用两个栈来保存target的前驱和后继, 而且栈顶的元素保存的是距离target最近的前驱和后继, 这样就可以每次取到距离target最近的值. 这种时间复杂度可以达到O(logn).
代码如下:
/**
* Definition for a binary tree node.
* struct TreeNode {
* int val;
* TreeNode *left;
* TreeNode *right;
* TreeNode(int x) : val(x), left(NULL), right(NULL) {}
* };
*/
class Solution {
public:
void initStack(TreeNode* root, double target)
{
while(root)
{
if(root->val <= target)
{
pre.push(root);
root = root->right;
}
else
{
succ.push(root);
root = root->left;
}
}
}
void getPre()
{
auto node = pre.top();
pre.pop();
if(node->left)
{
pre.push(node = node->left);
while(node->right)
pre.push(node = node->right);
}
}
void getSucc()
{
auto node = succ.top();
succ.pop();
if(node->right)
{
succ.push(node = node->right);
while(node->left)
succ.push(node = node->left);
}
}
vector closestKValues(TreeNode* root, double target, int k) {
if(!root) return {};
initStack(root, target);
while(k-- > 0)
{
if(pre.empty()||(!succ.empty() && fabs(pre.top()->val-target) > fabs(succ.top()->val-target)))
{
ans.push_back(succ.top()->val);
getSucc();
}
else
{
ans.push_back(pre.top()->val);
getPre();
}
}
return ans;
}
private:
stack pre, succ;
vector ans;
};