Binary Search Tree Iterator

Problem

Design an iterator over a binary search tree with the following rules:

Elements are visited in ascending order (i.e. an in-order traversal) next() and hasNext() queries run in O(1) time in average. Have you met this question in a real interview? Yes Example For the following binary search tree, in-order traversal by using iterator is [1, 6, 10, 11, 12]

   10
 /    \
1      11
 \       \
  6       12

思路

首先注意是binary search tree, 所以所有左children小于root,右children大于root

next()要求是不断能找到下一个最小的数,在BST中,最小的node是the most left node 所以用一个stack来recursively push root所有的left node, pop()得到最小的数

注意 如果当前pop出来的数有right child,就要把下一个smallest node push到stack上,也就是将right child的所有left children push进去

hasNext()只要check stack是否为空就好

Solution


"""
Definition of TreeNode:
class TreeNode:
    def __init__(self, val):
        self.val = val
        self.left, self.right = None, None

Example of iterate a tree:
iterator = Solution(root)
while iterator.hasNext():
    node = iterator.next()
    do something for node 
"""
class Solution:
    #@param root: The root of binary tree.
    def __init__(self, root):
        # write your code here
        self.stack = []
        self.pushAllLeft(root)

    #@return: True if there has next node, or false
    def hasNext(self):
        # write your code here
        return self.stack

    #@return: return next node
    def next(self):
        #write your code here
        node = self.stack.pop()
        if node.right:
            self.pushAllLeft(node.right)
        return node

    def pushAllLeft(node):
        while node:
            self.stack.append(node)
            node = node.left