My C++ code using getPredecessor and getSuccessor


  • 0
    O

    It can terminate earlier if both the predecessor and the successor are worse than the existing top-k.
    Please share your comments.

    class Solution {
        stack<TreeNode *> predSt;
        stack<TreeNode *> succSt;
        
        TreeNode *getPredecessor(TreeNode *node)
        {
            if (node->left)
            {
                predSt.push(node);
                node = node->left;
                while (node->right)
                {
                    predSt.push(node);
                    node = node->right;
                }
                return node;
            }
            else
            {
                while (!predSt.empty())
                {
                    TreeNode *parent = predSt.top();
                    predSt.pop();
                    if (parent->right == node)
                        return parent;
                    node = parent;
                }
                return NULL;
            }
        }
        
        TreeNode *getSuccessor(TreeNode *node)
        {
            if (node->right)
            {
                succSt.push(node);
                node = node->right;
                while (node->left)
                {
                    succSt.push(node);
                    node = node->left;
                }
                return node;
            }
            else
            {
                while (!succSt.empty())
                {
                    TreeNode *parent = succSt.top();
                    succSt.pop();
                    if (parent->left == node)
                        return parent;
                    node = parent;
                }
                return NULL;
            }
        }
        
    public:
        vector<int> closestKValues(TreeNode* root, double target, int k) {
            priority_queue<pair<double, int>> heap;
            heap.push(make_pair(fabs(target - root->val), root->val));
            TreeNode *pred = getPredecessor(root);
            TreeNode *succ = getSuccessor(root);
            while (true)
            {
                if (!pred && !succ) //traversed all nodes
                    break;
                
                double diffPred = (pred == NULL) ? numeric_limits<double>::max() : fabs(target - pred->val);
                double diffSucc = (succ == NULL) ? numeric_limits<double>::max() : fabs(target - succ->val);
                
                //the unseen nodes cannot be better than the existing top-k ones, stop traversal
                if (heap.size() == k && heap.top().first <= min(diffPred, diffSucc))
                    break;
                
                if (diffPred < diffSucc)
                {
                    heap.push(make_pair(fabs(diffPred), pred->val));
                    pred = getPredecessor(pred);
                }
                else
                {
                    heap.push(make_pair(fabs(diffSucc), succ->val));
                    succ = getSuccessor(succ);
                }
                if (heap.size() > k)
                    heap.pop();
            }
            vector<int> res;
            while (!heap.empty())
            {
                const pair<double, int> &top = heap.top();
                res.push_back(top.second);
                heap.pop();
            }
            return res;
        }
    };

Log in to reply
 

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