[Top Interview 150] 173. Binary Search Tree Iterator

Solution)

class BSTIterator:
    def __init__(self, root: Optional[TreeNode]):
        # arr stores h nodes (h = height of the tree)
        self.arr = []
        
        # Helper function adds nodes in root => leftmost leaf branch
        self.findNextNodes(root)


    def next(self) -> int:
        # nextNode is the smallest number node
        nextNode = self.arr.pop()

        # If nextNode has a right child, we add nodes from
        # right child => leftmost child to arr
        if nextNode.right:
            self.findNextNodes(nextNode.right)
        
        return nextNode.val

    
    # If there is a node in arr, return True
    def hasNext(self) -> bool:
        if self.arr:
            return True

        return False
    

    # Helper function adds nodes from root => leftmost child
    # We can access the smallest number node and the second smallest number node
    def findNextNodes(self, node):
        while node:
            self.arr.append(node)
            node = node.left

 
Problem: LeetCode
 
*Time Complexity: O(1)
Space Complexity: O(h)
* Average of O(1) applies to next() and hasNext()
 

Explanation)

We need to find the smallest node (smallest node as in node with the smallest value).
 
As we find the smallest node, we'll add the nodes in its path from root => smallest node.
 
The reason we do this is we can,
1) get the smallest node
2) get the second smallest node
3) access the right subtree of that second smallest node
 

Adding nodes from 7 => 5 => 3

 

1. Helper Function)

To make our code more readable, we'll have a helper function.
 
This helper function findNextNodes() will add nodes from the current position to its leftmost child.
 
We do this to get to the smallest node in the branch.
 

class BSTIterator:
    # 1. Helper Function
    def findNextNodes(self, node):
        while node:
            self.arr.append(node)
            node = node.left

 

2. Initialization)

Variable arr will be our stack that saves nodes using findNextNodes().
 
In our initialization, we'll execute findNextNodes() to save nodes from root to leftmost child. 
Now that findNextNodes(root) is executed, we have the smallest node in arr.
 

class BSTIterator:
    # 1. Helper Function
    def findNextNodes(self, node):
        while node:
            self.arr.append(node)
            node = node.left


    # 2. Initialization
    def __init__(self, root: Optional[TreeNode]):
        self.arr = []
        self.findNextNodes(root)

 

3. next)

In short, next() will return the smallest integer to the biggest integer in a tree.
 
From our initialization, we have the smallest node in arr's -1st position.
So, we'll pop that node and save that in a variable.
 
Now we need to check if this smallest node has children.
 
Because we already know that the smallest node is the leftmost child, we also know that there is no left child.
Then, we gotta check its right child.
 
If the smallest node has a right child, we have to execute findNextNodes() with the right child again, so that we can get the next smallest child.
 

If the current node has a right child

 
After executing findNextNodes(), don't forget to return the current node's value.
 

class BSTIterator:
    # 1. Helper Function

    # 2. Initialization
    def __init__(self, root: Optional[TreeNode]):
        self.arr = []
        self.findNextNodes(root)


    # 3. next
    def next(self) -> int:
        nextNode = self.arr.pop()

        if nextNode.right:
            self.findNextNodes(nextNode.right)
        
        return nextNode.val

 

4. hasNext)

This function is very simple.
 
If arr is not empty, there are still nodes to return.
If arr is empty, there is no nodes to return.
 

class BSTIterator:
    # 1. Helper Function

    # 2. Initialization

    # 3. next
    def next(self) -> int:
        nextNode = self.arr.pop()

        if nextNode.right:
            self.findNextNodes(nextNode.right)
        
        return nextNode.val

    
    # 4. hasNext
    def hasNext(self) -> bool:
        if self.arr:
            return True

        return False

---

 
 
I thought my next() function's time complexity is O(h), but apparently it's O(1) on average. I think I'm confused between amorteized complexity and average complexity. I'll make a post about it after studying it.