Link to the 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)?

Example

Input: [2, 1, 3], target = 1.5, k = 2, Output: [1, 2]

Idea

First do an inorder traversal, then find the k closest values by finding the pivot, and run the merge in merge sort.

Solution

class Solution {
private:
    void inorderTraversal(TreeNode *root, vector<int> &values) {
        if (root) {
            inorderTraversal(root->left, values);
            values.push_back(root->val);
            inorderTraversal(root->right, values);
        }
    }

public:
    vector<int> closestKValues(TreeNode* root, double target, int k) {
        vector<int> values;
        inorderTraversal(root, values);
        vector<int> rtn;
        int n = values.size();
        int right = 0;
        while (right < n && values[right] < target) right++;
        int left = right - 1;
        while (rtn.size() < k) {
            if (left >= 0 && right < n) {
                if (abs((double) values[left] - target) < abs((double) values[right] - target)) {
                    rtn.push_back(values[left--]);
                } else {
                    rtn.push_back(values[right++]);
                }
            } else if (left >= 0) {
                rtn.push_back(values[left--]);
            } else {
                rtn.push_back(values[right++]);
            }
        }
        return rtn;
    }
};

68 / 68 test cases passed.
Runtime: 13 ms