C++ solution, using priority queue, O(n) time, O(k) space


  • 1
    L
    struct compare {
       bool operator()(pair<double, int> &a, pair<double, int> &b) { return a.first < b.first;}
    };
    
    typedef priority_queue<pair<double, int>, vector<pair<double, int>>, compare> MQ;
    vector<int> closestKValues(TreeNode* root, double target, int k) {
        MQ max_heap;
        vector<int> res;
        if (!root) return res;
        inorder(root, max_heap, target, k);
        while (!max_heap.empty()) { res.push_back(max_heap.top().second); max_heap.pop(); }
        return res;
    }
    void inorder(TreeNode *root, MQ &que, double target, int k) {
        if (!root) return;
        inorder(root->left, que, target, k);
        if (que.size() < k || abs(root->val - target) < que.top().first) {
            que.emplace(abs(root->val - target), root->val);
            if (que.size() > k) que.pop();
        }
        inorder(root->right, que, target, k);
    }

  • 0
    C

    First time for me to see
    typedef priority_queue<pair<double, int>, vector<pair<double, int>>, compare> MQ;

    Though my solution is quite slow.

    /**
     * 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) {
            priority_queue<double> distQ;
            unordered_map<double, vector<int>> dist2Val;
            
            traverse(root,target,distQ,dist2Val);
            vector<int> res;
            
            for (int i=0; i<k; )
            {
                double key = distQ.top();
                distQ.pop();
                vector<int> vals=dist2Val[key];
                int numVals = vals.size();
                while(numVals>1)
                {
                    distQ.pop();
                    numVals--;
                }
                for (auto val : vals)
                {
                    res.push_back(val);
                    i++;
                    if (i>=k)
                      break;
                }
            }
            
            return res;
        }
        
        void traverse(TreeNode *cur, double target, priority_queue<double> & distQ, unordered_map<double, vector<int>> &dist2Val)
        {
            if (cur==NULL)
               return;
               
            double dist=abs(target-cur->val);
            
            distQ.push(-dist);
            dist2Val[-dist].push_back(cur->val);
            
            traverse(cur->left, target, distQ, dist2Val);
            traverse(cur->right, target, distQ, dist2Val);
        }
    };

  • 0
    Z

    Hi, I have the same idea, but why the time complexity is O(n)? Thanks!


Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.