https://leetcode.com/problems/binary-search-tree-iterator/
Binary Search Tree Iterator
Implement an iterator over a binary search tree (BST). Your iterator will be initialized with the root node of a BST.
Calling next() will return the next smallest number in the BST.
Note: next() and hasNext() should run in average O(1) time and uses O(h) memory, where h is the height of the tree.
Credits:
Special thanks to @ts for adding this problem and creating all test cases.
Code:
# Definition for a binary tree node
# class TreeNode(object):
# def __init__(self, x):
# self.val = x
# self.left = None
# self.right = None
class BSTIterator(object):
def __init__(self, root):
"""
:type root: TreeNode
"""
self.stack = []
self.pushAll(root)
print [s.val for s in self.stack]
def hasNext(self):
"""
:rtype: bool
"""
return self.stack
def next(self):
"""
:rtype: int
"""
n = self.stack.pop()
self.pushAll(n.right)
return n.val
def pushAll(self, node):
while node:
self.stack += [node]
node = node.left
# Your BSTIterator will be called like this:
# i, v = BSTIterator(root), []
# while i.hasNext(): v.append(i.next())
Idea
First, load `self.stack` from root to the leftmost leaf. This requires O(log n) memory and time. Then, when you call `next()`, you pop the top element from `self.stack`, which is guaranteed to be smallest. Before you return this, you also need to load this element’s right child’s path to its leftmost child into `self.stack`. To do so, you call `self.pushAll(node.right)`.