Concise C++ sol using 2 stacks


  • 0
    Z
    class Solution {
    public:
        vector<int> closestKValues(TreeNode* root, double target, int k) {
            stack<TreeNode *> sts, stb;
            vector<int> ret;
            while (root) {
                if (root->val < target) {
                    sts.push(root);
                    root = root->right;
                } else {
                    stb.push(root);
                    root = root->left;
                }
            }
            while (k-- != 0) {
                if (stb.empty() || (!sts.empty() && target - sts.top()->val < stb.top()->val - target)) {
                    ret.push_back(sts.top()->val);
                    nextNodeInSTS(sts);
                } else {
                    ret.push_back(stb.top()->val);
                    nextNodeInSTB(stb);
                }
            }
            return ret;
        }
    private:
        void nextNodeInSTB(stack<TreeNode *> &stb) {
            TreeNode *tmp = stb.top()->right;
            stb.pop();
            while (tmp) {
                stb.push(tmp);
                tmp = tmp->left;
            }
        }
    
        void nextNodeInSTS(stack<TreeNode *> &sts) {
            TreeNode *tmp = sts.top()->left;
            sts.pop();
            while (tmp) {
                sts.push(tmp);
                tmp = tmp->right;
            }
        }
    };
    

Log in to reply
 

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