Closest Binary Search Tree Value II.cpp

思路就是对predecessor和successor分别建立iterator，然后根据情况决定这一次是选predecessor的next还是successor的next，直到我们取到k个值位置。建立Binary Tree Iterator的思路可以参考这一题。这一题是类似的，我们以successor为例，区别就是我们不会把小于target的数push进stack，因为我们发现当前节点curr的值小于target的话，那么整个左子树我们就不需要考察了，直接去右子树继续构建stack。如果当前节点curr的值大于等于target，那么我们存入stack，方便之后遍历其右子树。显而易见，简历stack的时间复杂度为O(log n)，有了stack之后，我们就可以按照iterator的方法每次取一个值了。注意handle 和target值相等的node的时候，只需要predecessor和successor iterator其中一个handle就可以了。时间复杂度等于建stack + using iterator to find k elements = O(log n + k)，代码如下：

	/**
	* 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<int> closestKValues(TreeNode* root, double target, int k) {
	stack<TreeNode*> pre, suc;
	initializeSucStack(suc, target, root);
	initializePreStack(pre, target, root);
	vector<int> res;
	while(k-- > 0)
	{
	if(pre.empty())
	res.push_back(getNextSuc(suc));
	else if(suc.empty())
	res.push_back(getNextPre(pre));
	else
	{
	double diffPre = abs(target - pre.top()->val), diffSuc = abs(target - suc.top()->val);
	if(diffPre <= diffSuc)res.push_back(getNextPre(pre));
	else res.push_back(getNextSuc(suc));
	}
	}
	return res;
	}
	private:
	void initializeSucStack(stack<TreeNode*>& st, double target, TreeNode* node)
	{
	while(node)
	{
	if(node->val < target)
	node = node->right;
	else
	{
	st.push(node);
	node = node->left;
	}
	}
	}
	void initializePreStack(stack<TreeNode*>& st, double target, TreeNode* node)
	{
	while(node)
	{
	if(node->val >= target)
	node = node->left;
	else
	{
	st.push(node);
	node = node->right;
	}
	}
	}
	int getNextSuc(stack<TreeNode*>& st)
	{
	if(st.empty())return -1;
	auto res = st.top();
	st.pop();
	auto curr = res->right;
	while(curr)
	{
	st.push(curr);
	curr = curr->left;
	}
	return res->val;
	}
	int getNextPre(stack<TreeNode*>& st)
	{
	if(st.empty())return -1;
	auto res = st.top();
	st.pop();
	auto curr = res->left;
	while(curr)
	{
	st.push(curr);
	curr = curr->right;
	}
	return res->val;
	}
	};

view raw Closest Binary Search Tree Value II.cpp hosted with ❤ by GitHub