173. Binary Search Tree Iterator - cocoder39/coco39_LC GitHub Wiki

173. Binary Search Tree Iterator

an intuitive idea is to do in-order traversal to collect nodes in ascending order. For in-order traversal, we can either use recursion or iterative solution by leveraging stack. The optimization here is to modify the iterative solution.

Basic in-order traversal

class BSTIterator:

    def __init__(self, root: Optional[TreeNode]):
        self.idx = 0
        self.sorted_array = []
        stack = []
        while root or stack:
            if root:
                stack.append(root)
                root = root.left
            else:
                node = stack.pop()
                self.sorted_array.append(node)
                root = node.right

    def next(self) -> int:
        res = self.sorted_array[self.idx]
        self.idx += 1
        return res.val

    def hasNext(self) -> bool:
        return self.idx < len(self.sorted_array)

Optimization: split the original iterative solution to several parts to be distributed in constructor and next(). The smallest node at current moment is always the top of the stack. O(hight) space complexity

class BSTIterator:

    def __init__(self, root: Optional[TreeNode]):
        self.stack = []
        while root:
            self.stack.append(root)
            root = root.left

    def next(self) -> int:
        root = self.stack.pop()
        res = root.val
        root = root.right
        while root:
            self.stack.append(root)
            root = root.left
        return res

    def hasNext(self) -> bool:
        return len(self.stack) > 0

==========================================================

O(1) average time and O(height) space

class BSTIterator {
public:
    BSTIterator(TreeNode *root) {
        stackPush(root);
    }

    /** @return whether we have a next smallest number */
    bool hasNext() {
        return ! s.empty();
    }

    /** @return the next smallest number */
    int next() {
        TreeNode* node = s.top();
        s.pop();
        int res = node->val;
        stackPush(node->right);
        return res;
    }
private:
    stack<TreeNode*> s;
    void stackPush(TreeNode* node) {
        while (node) {
            s.push(node);
            node = node->left;
        }
    }
};

pre-order iterator: average O(1) time and O(h) memeory

class BSTIterator {
public:
    BSTIterator(TreeNode *root) {
        this->root = root;
    }

    /** @return whether we have a next smallest number */
    bool hasNext() {
        return root != nullptr;
    }

    /** @return the next smallest number */
    int next() {
        stk.push(root);
        int res = root->val;
        root = root->left;
        while (!root && !stk.empty()) {
            TreeNode* node = stk.top();
            stk.pop();
            root = node->right;
        }
        return res;
    }
private:
    TreeNode* root;
    stack<TreeNode*> stk;
};

post order iterator

⚠️ **GitHub.com Fallback** ⚠️