C++ 20ms solution using a TreeIterator


  • 1
    K
    class InorderIterator : public iterator<forward_iterator_tag, TreeNode> {
        private:
            TreeNode* cur;
            stack<TreeNode*> stk;
            bool valid;
            
            void advance() {
                TreeNode* tmp;
                if (!stk.empty()) {
                    cur = stk.top();
                    stk.pop();
                    tmp = cur->right;
                    while (tmp) {
                        stk.push(tmp);
                        tmp = tmp->left;
                    }
                } else {
                    valid = false;
                }
            }
            
        public:
            explicit InorderIterator(TreeNode* root) {
                cur = root;
                while (cur) {
                    stk.push(cur);
                    cur = cur->left;
                }
                advance();
                valid = true;
            }
            
            InorderIterator& operator++() {
                advance();
                return *this;
            }
            
            TreeNode& operator*() {
                return *cur;
            }
            
            TreeNode* operator->() {
                return cur;
            }
            
            bool hasNext() {
                return !stk.empty();
            }
            
            bool isValid() {
                return valid;
            }
    };
    
    
    vector<int> closestKValues(TreeNode* root, double target, int k) {
        InorderIterator it1(root), it2(root);
        double avg = 0.0, minavg;
        int minleft = 0, i = 0;
        for (; i < k; i++) {
            avg += it2->val;
            ++it2;
        }
        minavg = abs(avg-target*k);
        i = 0;
        while (it2.isValid()) {
            avg += it2->val;
            avg -= it1->val;
            ++it2;
            ++it1;
            i++;
            if (abs(avg-target*k) < minavg) {
                minavg = abs(avg-target*k);
                minleft = i;
            }
        }
        
        vector<int> result;
        InorderIterator it3(root);
        for (int i = 0; i < minleft; i++) ++it3;
        for (int i = 0; i < k; i++) {
            result.push_back(it3->val);
            ++it3;
        }
        return result;
    }

Log in to reply
 

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