天天看點

LeetCode 173. Binary Search Tree IteratorLeetCode 173

LeetCode 173

最簡單的方法就是,先把樹flatten中序周遊成一個list,然後在next中逐個傳回。

但是note中有一個條件,空間複雜度隻能是O(h),那麼我們不能把整個樹轉為list存下來。

第一個想到的就是利用stack來做中序周遊,中序周遊基于stack的方法是:

def inOrder(root): 
    current = root  
    stack = [] 
     
    while True: 
        if current is not None: 
            stack.append(current)         
            current = current.left             
        elif(stack): 
            current = stack.pop() 
            print(current.data, end=" ") 
            current = current.right  
        else: 
            break
           

基于這個,我們可以把iterator實作成:

class BSTIterator:

    def __init__(self, root: TreeNode):        
        node = root
        self.stack = []
        while node:
            self.stack.append(node)
            node = node.left

    def next(self) -> int:
        """
        @return the next smallest number
        """
        result = self.stack.pop()
        if result.right:
            node = result.right
            while node:
                self.stack.append(node)
                node = node.left

        return result.val

    def hasNext(self) -> bool:
        """
        @return whether we have a next smallest number
        """
        return len(self.stack) > 0