Problem statement

https://leetcode.com/problems/binary-search-tree-iterator-ii/

Solution 1

One way to solve is just create traversal and then use created list.

Complexity

Time complexity if init is O(n), all others are O(1).

Code

class BSTIterator:
    def __init__(self, root):
        def dfs(node):
            if not node: return []
            return dfs(node.left) + [node.val] + dfs(node.right)
        self.arr = dfs(root)
        self.size = len(self.arr)
        self.idx = - 1

    def hasNext(self):
        return self.idx < self.size - 1

    def next(self):
        self.idx += 1
        return self.arr[self.idx]
        
    def hasPrev(self):
        return self.idx > 0

    def prev(self):
        self.idx -= 1
        return self.arr[self.idx]

Solution 2

No need to traverse all tree: we can do it only when we need. We can use Iterative Inorder Traversal.

Complexity

Time complexity of all operations is O(1) except next. For next amorized complexity is O(1). Space complexity is O(n) due to self.arr.

Code

class BSTIterator:
    def __init__(self, root):
        self.last = root
        self.stack, self.arr = [], []
        self.idx = -1

    def hasNext(self):
        return self.stack or self.last or self.idx < len(self.arr) - 1

    def next(self):
        self.idx += 1
        
        if self.idx == len(self.arr):
            while self.last:
                self.stack.append(self.last)
                self.last = self.last.left
            curr = self.stack.pop()
            self.last = curr.right
            self.arr.append(curr.val)
            
        return self.arr[self.idx]

    def hasPrev(self):
        return self.idx > 0
    
    def prev(self):
        self.idx -= 1
        return self.arr[self.idx]