272. Closest Binary Search Tree Value II

• Difficulty: Hard

• Topics: Stack, Tree

• Similar Questions:

  • Binary Tree Inorder Traversal

  • Closest Binary Search Tree Value

Problem:

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.

Example:

Input: root = [4,2,5,1,3], target = 3.714286, and k = 2

    4
   / \
  2   5
 / \
1   3

Output: [4,3]

Follow up:
Assume that the BST is balanced, could you solve it in less than O(n) runtime (where n = total nodes)?

Solutions:

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */

class Iterator{
    public:
    Iterator(stack<TreeNode*> stk, bool forward): stk_(stk), forward_(forward) {

    }

    int next() {
        TreeNode* node = stk_.top(); stk_.pop();
        int ret = node->val;
        node = forward_ ? node->right : node->left;
        while (node) {
            stk_.push(node);
            node = forward_ ? node->left : node->right;
        }
        return ret;
    }

    // remember to has peek()
    int peek() {
        return stk_.top()->val;
    }

    bool hasNext() {
        return !stk_.empty();
    }

    private:
        bool forward_;
        stack<TreeNode*> stk_;
};

class Solution {
public:
    vector<int> closestKValues(TreeNode* root, double target, int k) {
        if (k == 0) return {};
        if (root == NULL)   return {};
        stack<TreeNode*> fstk;
        stack<TreeNode*> bstk;
        findHelper(root, target, fstk, bstk);

        Iterator forwardIter = Iterator(fstk, true);
        Iterator backIter = Iterator(bstk, false);

        if (forwardIter.peek() <= target) {
            forwardIter.next();
        } else {
            backIter.next();
        }


        vector<int> ret;
        for (int i = 0; i < k; ++i) {
            if (!forwardIter.hasNext()) {
                ret.push_back(backIter.next());
                continue;
            }

            if (!backIter.hasNext()) {
                ret.push_back(forwardIter.next());
                continue;
            }

            int forwardVal = forwardIter.peek();
            int backVal = backIter.peek();

            if (abs(forwardVal - target) <= abs(backVal - target)) {
                ret.push_back(forwardVal);
                if (forwardIter.hasNext()) {
                    forwardIter.next();
                }
            } else {
                ret.push_back(backVal);
                if (backIter.hasNext()) {
                    backIter.next();
                }
            }
        }
        return ret;
    }

    // differciate fstk and bstk!
    void findHelper(TreeNode* root, double target, stack<TreeNode*>& fstk, stack<TreeNode*>& bstk) {   
        if (double(root->val) == target) {
            fstk.push(root);
            bstk.push(root);
            return;
        } 

        if (double(root->val) > target) {
            if (root->left) {
                fstk.push(root);
                findHelper(root->left, target, fstk, bstk);
                return;
            }

            fstk.push(root);
            bstk.push(root);
            return;
        } 
        if (double(root->val) < target) {
            if (root->right) {
                bstk.push(root);
                findHelper(root->right, target, fstk, bstk);
                return;
            }
            fstk.push(root);
            bstk.push(root);
            return;
        }
    }
};