C++ Solution Long but Clear O(log(n) + k)


  • 0
    X
    class Solution {
        TreeNode * smaller, * greater;
        stack<TreeNode *> prevStk, nextStk;
    public:
        vector<int> closestKValues(TreeNode* root, double target, int k) {
            TreeNode * p = closestValue(root, target);
            smaller = p; // Set pointers
            greater = p;
            while(root){ // Set Stacks
                if(p->val == root->val)
                    break;
                else{
                    prevStk.push(root);
                    nextStk.push(root);
                    if(root->val < p->val) root = root->right;
                    else root = root->left;
                }
            }
            vector<int> res;
            res.push_back(p->val);
            k--;
            next();
            prev();
            while(k > 0){
                k--;
                if(!smaller && !greater) break;
                else if( !greater || ( smaller && (target - (double)smaller->val <= (double)greater->val - target))){
                    res.push_back(smaller->val);
                    prev();
                }
                else if( smaller == NULL || (greater && (target - (double)smaller->val >= (double)greater->val - target))){
                    res.push_back(greater->val);
                    next();
                }
            }
            return res;
        }
        
        TreeNode * closestValue(TreeNode* root, double target) {
            if(target > INT_MAX) target = (double)INT_MAX;
            if(target < INT_MIN) target = (double)INT_MIN;
            TreeNode * node = root;
            while(root){
                if(abs(target - root->val) < abs(target - node->val)){
                    node = root;
                }
                if(target > root->val)
                    root = root->right;
                else if(target < root->val)
                    root = root->left;
                else return root;
            }
            return node;
        }
        
        void prev(){
            if(smaller->left){
                smaller = smaller->left;
                while(smaller->right) {prevStk.push(smaller); smaller = smaller->right;}
            }
            else{
                while(!prevStk.empty()){
                    if(prevStk.top()->val < smaller->val){
                        smaller = prevStk.top();
                        prevStk.pop();
                        return;
                    }
                    else prevStk.pop();
                }
                smaller = NULL;
            }
        }
        
        void next(){
            if(greater->right){
                greater = greater->right;
                while(greater->left) {nextStk.push(greater);greater = greater->left;}
            }
            else{
                while(!nextStk.empty()){
                    if(nextStk.top()->val > greater->val){
                        greater = nextStk.top();
                        nextStk.pop();
                        return;
                    }
                    else nextStk.pop();
                }
                greater = NULL;
            }
        }
    };

  • 0

    It's O(k log n) again. (the complexity of prev() and next() are both log n, and they will be called O(k) times)


Log in to reply
 

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