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.

二叉查找树的一道设计题。但是其实思路比较难想。

首先需要均摊时间复杂度为O(1)，使用O(h)复杂度度，其实很容易想到使用栈来做，但是如何来做比较难想。首先返回一个递增的序列，还是中序遍历。但是直接先用中序遍历存储所有结点，之后再返回的空间复杂度为O(n),不符合题目O(h)的空间复杂度的要求。所以该题的做法应该在next中给栈增加元素。设置一个self.cur来表明当前指向的结点。但是该结点不一定是当前需要返回的元素。其实是要返回以该结点为根的子树中的最左结点。比如开始的时候 self.cur = root。但是需要寻找以root为根的最左结点。为了以后的回溯方便，我们需要在找寻最左结点时保存遍历到的结点。在遍历到最左结点之后，返回该结点。则为了下一步的next准备我们把self.cur设置为返回结点的右孩子。这符合中序遍历的做法。如果该右孩子不存在，则利用栈中元素进行回溯（弹出栈），这个回溯到的点是直接需要返回的值,时间复杂度O(1)。如果该右孩子为叶子结点，则直接返回该叶子结点，时间复杂度也为O(1)。最坏的情况是该右孩子有子树，则需要找到该子树中的最左结点。这个的时间复杂度不是O(1)。但是每个结点一次被压栈，一次被弹出，其实平均的时间复杂度为O(1)。代码如下：

class BSTIterator(object):
    def __init__(self, root):
        """
        :type root: TreeNode
        """
        self.stack = []
        self.cur = root

    def hasNext(self):
        """
        :rtype: bool
        """
        return len(self.stack)>0 or self.cur != None
    
    def next(self):
        """
        :rtype: int
        """
        while self.cur: 
            self.stack.append(self.cur)
            self.cur = self.cur.left
        self.cur = self.stack.pop() #self.cur now is the left most node.
        node = self.cur
        self.cur = self.cur.right
        return node.val

 简单做法：

其实就是找二叉查找树的后继

# 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 = []
        cur = root
        while cur:
            self.stack.append(cur)
            cur = cur.left

    def hasNext(self):
        """
        :rtype: bool
        """
        return len(self.stack)>0
    
    def next(self):
        """
        :rtype: int
        """
        node = self.stack.pop()
        cur = node.right
        while cur:
            self.stack.append(cur)
            cur = cur.left
        return node.val
        
# Your BSTIterator will be called like this:
# i, v = BSTIterator(root), []
# while i.hasNext(): v.append(i.next())

这题的followup 如果结点有属性parent记录父结点，要求用O(1)的空间解：

class BSTIteratorN(object):
    def __init__(self, root):
        """
        :type root: TreeNode
        """
        self.prev = None
        self.cur = root
        while self.cur and self.cur.left:
            self.cur = self.cur.left

    def hasNext(self):
        """
        :rtype: bool
        """
        return self.cur != None

    def next(self):
        """
        :rtype: int
        """
        node = self.cur
        if self.prev == None:
            self.cur = self.cur.parent
        else:
            if self.prev == self.cur.parent:  #最后从这里结束遍历
                self.cur = self.cur.parent.parent
            #注意这里不能直接验证是否和左子结点相等，存在右子树完全走完，向上跳两级父结点的情况。此时也需要遍历当前结点的右子树。但是prev不是当前结点的左结点
            elif self.prev != self.cur.right:
                self.cur = self.cur.right
                while self.cur and self.cur.left:
                    self.cur = self.cur.left
        self.prev = node
        return node.val