1 題目描述

簡單來說，這題是要實作一個Binary Search Tree的iterator，這個iterator可以讓你用next()的方式來取得下一個最小的值。然後使用 hasNext() 來判斷是否還有下一個最小值。

注意：題目有要求，空間複雜度必須是O(h)，h是樹的高度。

2 解法

當你看到 next() 目的是在取得『最小值』時。因為BST的特性是，任意節點的左子樹不空，那左子樹上所有節點的值均小於他的根節點的值。因此我們可以使用 inorder traversal 的方式來解這題，因為 inorder traversal 是一種由小到大(左-中-右)的方式遍歷BST。詳細關於inorder得實作可以參考這篇LeetCode 課前預習 - 掌握 BFS 與 DFS 指南。

2.1 解法1: Inorder Traversal

這邊我們可以使用stack來幫助我們做inorder traversal，然後每次呼叫next()的時候，我們就可以取得最小的值，在 hasNext() 的時候，我們只要判斷stack是否為空就可以了。

class BSTIterator:

    def __init__(self, root: Optional[TreeNode]):
        self.stack = []
        self._inorder(root)

    def _inorder(self, node):
        if node is None:
            return
        self._inorder(node.right)   # 因為 FILO，所以先放右邊（大的）比較晚出來
        self.stack.append(node)
        self._inorder(node.left)    # 再放左邊（小的）比較早出來

    def next(self) -> int:
        return self.stack.pop().val

    def hasNext(self) -> bool:
        return len(self.stack) > 0

但是上面的解法，空間複雜度是O(n)，不符合題目要求，上面的程式問題在於 _inorder 的時候，把所有節點都放進stack裡面，這樣空間複雜度就是O(n)。

2.2 解法2: Iterator with Stack

在解法1中，我們把所有的節點都保存起來了，但其實沒必要一次性保存所有節點，我們應該要控制stack塞入節點的時機，這樣就可以達到O(h)的空間複雜度。具體該怎麼做，我們來釐清一下：

第一種情境

當我們初始化BSTIterator的時候，在呼叫next()時，我們希望回傳最小值，而題目要求空間複雜度是O(h)，所以我們就遍歷到底，直到取到最小值，這個遍歷的過程中，我們可以把遍歷的節點都放進stack裡面，這樣就可以達到O(h)的空間複雜度。所以第一種狀況的程式碼大概是這樣：

1
2
3

while cur:
    stack.append(cur)
    cur = cur.left

第二種情境

接下來我們要考慮的是，當我們呼叫next()的時候，具體要考慮哪些事情，這裡分成兩個部分：

當我們呼叫next()的時候，我們要回傳最小值，也就是stack的最上面的節點。
當我們取出最小值後，指標往上移動，因為下次取next()時也是從stack中pop出來，但是我們的stack還沒有放任何右子樹的節點，因此每次pop出來後，我們就要檢查，right是否有值，如果有也要將右子樹遍歷到最小值的所有節點壓到stack中，這樣下次呼叫next()的時候，就可以取得下一個最小值。

res = stack.pop()
# 檢查是否有右子樹
if res.right:
    # 如果有，要把右子樹遍歷到最小值的所有節點壓到stack中，下次next()才會取得右子樹的最小值
    cur = res.right
    while cur:
        stack.append(cur)
        cur = cur.left
return res.val

那基本上我們就可以實作這個iterator了，我們把情境1的程式碼包裝成_leftmost_inorder()，最後程式碼如下：

class BSTIterator:
    def __init__(self, root: TreeNode):
        self.stack = []
        self._leftmost_inorder(root)

    def _leftmost_inorder(self, root):
        while root:
            self.stack.append(root)
            root = root.left

    def next(self) -> int:
        topmost_node = self.stack.pop()
        if topmost_node.right:
            self._leftmost_inorder(topmost_node.right)
        return topmost_node.val

    def hasNext(self) -> bool:
        return len(self.stack) > 0