c++, 16ms, merge k smaller and k greater, using stack


  • 0

    Well, the closestKSmallerValues() and closestKGreaterValues() are almost the same, they can be combined together if using an extra option parameter. However, I just give both of them.

    class Solution {
    public:
    vector<int> closestKSmallerValues(TreeNode* root, double target, int k) {
        vector<int> res;
        if (root==NULL) return res;
        stack<TreeNode*> path;
        TreeNode* cur=root;
        int cnt=0;
        while(cur){
            if(cur->val>target){
                cur=cur->left;
            }else if(cur->val<=target){
                path.push(cur);
                cur=cur->right;
            }
        }
        while(!path.empty()){
            cur = path.top();path.pop();
            res.push_back(cur->val);
            cnt++;
            if (cnt==k) return res;
            cur = cur->left;
            while(cur){
                path.push(cur);
                cur=cur->right;
            }
        }
        return res;
    }
    
    vector<int> closestKGreaterValues(TreeNode* root, double target, int k) {
        vector<int> res;
        if (root==NULL) return res;
        stack<TreeNode*> path;
        TreeNode* cur=root;
        int cnt=0;
        while(cur){
            if(cur->val<target){
                cur=cur->right;
            }else if(cur->val>=target){
                path.push(cur);
                cur=cur->left;
            }
        }
        while(!path.empty()){
            cur = path.top();path.pop();
            res.push_back(cur->val);
            cnt++;
            if (cnt==k) break;
            cur = cur->right;
            while(cur){
                path.push(cur);
                cur=cur->left;
            }
        }
        return res; 
    }
    
    vector<int> closestKValues(TreeNode* root, double target, int k) {
        vector<int> leq= closestKSmallerValues(root,target,k);
        vector<int> geq= closestKGreaterValues(root,target,k);
        vector<int> res(k,0);
        int idx1=0,idx2=0,idx;
        if (!geq.empty() && geq[0]==target) idx2++;
        while(idx1<leq.size() && idx2<geq.size() && idx<k){
            if (target-leq[idx1]<geq[idx2]-target){
                res[idx++]=leq[idx1++];
            }else{
                res[idx++]=geq[idx2++];
            }
        }
        if (idx==k) return res;
        if (idx1<leq.size()){
            while(idx<k){
                res[idx++]=leq[idx1++];
            }
        }
        if (idx2<geq.size()){
            while(idx<k){
                res[idx++]=geq[idx2++];
            }
        }
        return res;
    }
    };

Log in to reply
 

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