C++ Solution with dual Stacks for Successor and Predecessor Iterator


  • 0
    K

    The similar idea as https://discuss.leetcode.com/topic/23151/o-logn-java-solution-with-two-stacks-following-hint
    But with some enhancement when building the stacks.
    Instead of using target directly, I use its flooring value (int)target; by using its floor, we are building the predecessor iterator that no greater than target (included)
    And for successor, it's no less than target (exclude)

    predecessor : [low, target]
    successor: (target, high]

    Therefore, we can consolidate to while loop into one

        while(cur != NULL) {
            if (cur->val > val) {
                stkHi.push(cur);
                cur = cur->left;
            } else {
                stkLo.push(cur);
                cur = cur->right;
            }
        }
    

    Plus, its all integer comparison, so make it run slighter faster.
    The final result building block, we can potentially use (long)(target + 0.5) to find the closest value to compare and eliminate the FP compare more.

    For the time complexity, since it's similar to an ordinal in-order traversal without recursive which is O(n) but if k << n, it should be much faster than O(n) and regardless it's in the lower or higher side of the tree.

    class Solution {
    public:
        vector<int> closestKValues(TreeNode* root, double target, int k) {
            TreeNode *cur = root;
            long val = (int)target;
            stack<TreeNode *> stkHi;
            stack<TreeNode *> stkLo;
    
            while(cur != NULL) {
                if (cur->val > val) {
                    stkHi.push(cur);
                    cur = cur->left;
                } else {
                    stkLo.push(cur);
                    cur = cur->right;
                }
            }
    
            vector<int> res;
            while(res.size() < k) {
                TreeNode *lo = stkLo.size() ? stkLo.top() : NULL;
                TreeNode *hi = stkHi.size() ? stkHi.top() : NULL;
    
                if ((hi == NULL) ||
                    ((lo != NULL) && (target - lo->val < hi->val - target))) {
    
                    res.push_back(lo->val);
                    stkLo.pop();
                    lo = lo->left;
                    while(lo != NULL) {
                        stkLo.push(lo);
                        lo = lo->right;
                    }
                } else {
                    res.push_back(hi->val);
                    stkHi.pop();
                    hi = hi->right;
                    while(hi != NULL) {
                        stkHi.push(hi);
                        hi = hi->left;
                    }
                }
            }
    
            return res;
        }
    };
    

Log in to reply
 

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