272. 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:

• 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)?

void helper(TreeNode* node, double target, int k, priority_queue<pair<double, int>>& pq) {
    if (!node) return;
    pq.push(make_pair(abs(target - node->val), node->val));
    if (pq.size() > k) pq.pop();
    helper(node->left, target, k, pq);
    helper(node->right, target, k, pq);
}
vector<int> closestKValues(TreeNode* root, double target, int k) { // time: O(n*logk); space: O(n or k)
    vector<int> res;
    if (!root) return res;
    priority_queue<pair<double, int>> pq;
    helper(root, target, k, pq);
    while (!pq.empty()) {
        res.push_back(pq.top().second);
        pq.pop();
    }
    return res;
}

void getPredecessor(stack<TreeNode*>& pre) {
    TreeNode* t = pre.top(); pre.pop();
    // if (t->left) {
    //     pre.push(t->left);
    //     while (pre.top()->right) 
    //         pre.push(pre.top()->right);
    // }
    t = t->left;
    while (t) {
        pre.push(t);
        t = t->right;
    }
}
void getSuccessor(stack<TreeNode*>& suc) {
    TreeNode* t = suc.top(); suc.pop();
    // if (t->right) {
    //     suc.push(t->right);
    //     while (suc.top()->left)
    //         suc.push(suc.top()->left);
    // }
    t = t->right;
    while (t) {
        suc.push(t);
        t = t->left;
    }
}
vector<int> closestKValues(TreeNode* root, double target, int k) { // time: O(k*logn); space: O(logn)
    vector<int> res;
    stack<TreeNode*> pre, suc;
    TreeNode* cur = root;
    while (cur) {
        if (cur->val <= target) {
            pre.push(cur);
            cur = cur->right;
        } else {
            suc.push(cur);
            cur = cur->left;
        }
    }
    int n = k;
    while (n-- > 0) {
        if (suc.empty() || !pre.empty() && target - pre.top()->val < suc.top()->val - target) {
            res.push_back(pre.top()->val);
            getPredecessor(pre);
        } else {
            res.push_back(suc.top()->val);
            getSuccessor(suc);
        }
    }
    return res;
}

270. Closest Binary Search Tree Value