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