[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
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.
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.