Share O(klogn) time O(k + logn) memory C++ solution


  • 0
    H
    class Solution {
        stack<TreeNode*> parent_over;
        stack<TreeNode*> parent_under;
        TreeNode* root;
        TreeNode* curr;
        
        TreeNode* curr_succ;
        TreeNode* curr_prec;
    public:
        vector<int> closestKValues(TreeNode* _root, double target, int k) {
            if(!_root)   return {};
            if(k <= 0)  return {};
            
            root = _root;
            
            vector<int> result;
            find_closest(target);
            result.push_back(curr->val);
            --k;
            int last = curr->val;
            
            while(k--) {
                while(curr_prec && curr_prec->val == last) get_prec(); //we want distinct values
                while(curr_succ && curr_succ->val == last) get_succ(); 
                
                double error_prec = (curr_prec) ? fabs(target - curr_prec->val) : numeric_limits<double>::max();
                double error_succ = (curr_succ) ? fabs(target - curr_succ->val) : numeric_limits<double>::max();
                
                if(error_prec < error_succ) {
                    result.push_back(curr_prec->val);
                    last = curr_prec->val;
                    get_prec();
                } else {
                    result.push_back(curr_succ->val);
                    last = curr_succ->val;
                    get_succ();
                }
            }
            
            return result;
    
        }
        
        TreeNode* find_closest(double target) {
            TreeNode* temp = root;
            double best = numeric_limits<double>::max();
            int push_under_since_best = 0;
            int push_over_since_best = 0;
            
            while(temp) {
                double error = abs(target - temp->val);
                
                if(error < best) {
                    best = error;
                    curr = temp;
                    push_under_since_best = 0;
                    push_over_since_best = 0;
                }
                
                if(target < temp->val) {
                    if(!temp->left)  break;
                    parent_over.push(temp);
                    ++push_over_since_best;
                    temp = temp->left;
                } else if(target > temp->val) {
                    if(!temp->right)  break;
                    parent_under.push(temp);
                    ++push_under_since_best;
                    temp = temp->right;
                } else break;
            }
            
            while(push_under_since_best--)  parent_under.pop();
            while(push_over_since_best--)   parent_over.pop();
            
            curr_succ = curr;
            curr_prec = curr;
            get_succ();
            get_prec();
            
            return curr;
        }
        
        TreeNode* get_succ() {
            if(curr_succ->right) {
                curr_succ = curr_succ->right;
                while(curr_succ->left) {
                    parent_over.push(curr_succ);
                    curr_succ = curr_succ->left;
                }
            } else {
                if(parent_over.empty()) {
                    curr_succ = nullptr;
                    return nullptr;
                }
                curr_succ = parent_over.top();
                parent_over.pop();
            }
            
            return curr_succ;
        }
        
        TreeNode * get_prec() {
            if(curr_prec->left) {
                curr_prec = curr_prec->left;
                while(curr_prec->right) {
                    parent_under.push(curr_prec);
                    curr_prec = curr_prec->right;
                }
            } else {
                if(parent_under.empty()) {
                    curr_prec = nullptr;
                    return nullptr;
                }
                curr_prec = parent_under.top();
                parent_under.pop();
            }
            
            return curr_prec;
        }
    };

Log in to reply
 

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