173. Binary Search Tree Iterator

──────────────────────────────────────────────────────────────────────

173. Binary Search Tree Iterator

設計一個 BST 的迭代器，next() 依序回傳下一個最小的值、hasNext() 回答還有沒有下一個。

思路

第一個要看穿的是：「依序回傳 BST 的下一個最小值」就是「中序遍歷」。BST 的中序遍歷結果一定是遞增序列，所以這題等於「把中序遍歷拆成可以一次要一個的形式」。

有了這個翻譯，最直接的做法就出來了：在建構子裡先跑完整個中序，把結果存成一個 list 或 deque，next() 就是彈出一個。 這個版本一定要會寫，因為它把題目變成了「已知問題 + 一個佇列」。

但題目有一句 follow-up：

Could you implement next() and hasNext() to run in average $O(1)$ time and use $O(h)$ memory, where $h$ is the height of the tree?

$O(h)$ 才是這題真正在考的東西。 預先展開的版本是 $O(n)$ 空間 —— 樹有一百萬個節點時，建構子就要先把一百萬個值全部算出來存著，即使呼叫端只想拿前三個。

怎麼做到 O(h)

回想中序遍歷的 iterative 寫法（見 94. Binary Tree Inorder Traversal）：一路往左壓棧，彈出時處理，然後轉向右子樹。棧裡放的是「還沒被處理的祖先」，最多就是一條從根往下的路徑，所以是 $O(h)$。

這題的做法就是把那個迴圈拆開，暫停在每次彈出的地方：

• 建構子：從根一路往左壓棧
• next()：彈出棧頂（那就是目前最小的），然後把它右子樹的最左路徑壓進去
• hasNext()：棧是不是空的

這種「把一個遍歷過程拆成可以暫停 / 繼續的物件」叫做受控遞迴（controlled recursion），是設計類題目的常見手法。

解題方向

預先展開成佇列

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class BSTIterator:

    def __init__(self, root: Optional[TreeNode]):
        def inorder(node):
            res = []
            if not node:
                return []
            res.extend(inorder(node.left))
            res.append(node.val)
            res.extend(inorder(node.right))
            return res
        self.res = deque(inorder(root))

    def next(self) -> int:
        return self.res.popleft()


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


# Your BSTIterator object will be instantiated and called as such:
# obj = BSTIterator(root)
# param_1 = obj.next()
# param_2 = obj.hasNext()

正確，而且 next() / hasNext() 都是 $O(1)$ —— 但空間是 $O(n)$，沒有滿足 follow-up。

兩個小地方：hasNext 回傳的是 len(self.res)（一個 int）而不是 bool，Python 判斷真值時可以動，但和宣告的型別不符，寫 > 0 比較嚴謹。另外建構子裡的 res.extend(inorder(...)) 每層都在複製子樹的 list，在斜樹上會退化成 $O(n^2)$（94 那篇有說明）。

受控遞迴（$O(h)$ 空間）

class BSTIterator:

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

    def _push_left(self, node):
        while node:                    # 一路往左，沿途壓棧
            self.stack.append(node)
            node = node.left

    def next(self) -> int:
        node = self.stack.pop()        # 棧頂就是目前最小的
        self._push_left(node.right)    # 換去右子樹，一樣先壓到最左
        return node.val

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

_push_left 就是 94 那個 while cur: stack.append(cur); cur = cur.left 迴圈，只是被抽成一個方法，這樣建構子和 next() 都能用。

空間是 $O(h)$ 而不是 $O(n)$：stack 裡永遠只放「從當前位置一路往左」的那條路徑，所以平衡 BST 建構完只會有樹高那麼多個元素，只有左斜鏈才會退化成 n。對照預先展開的版本 —— 它建構完就佔了 n 個。

next() 是均攤 $O(1)$，不是最壞 $O(1)$：某一次呼叫可能要壓好幾層（最多 $O(h)$），但整個迭代過程中每個節點只會被壓入一次、彈出一次，所以 n 次呼叫總共 $O(n)$，平均下來是 $O(1)$。這個「均攤」要主動講，否則面試官會以為你沒注意到最壞情況。

補充

同一個中序骨架的題目：94. Binary Tree Inorder Traversal（本題的迴圈原型）、230. Kth Smallest Element in a BST（中序走到第 k 個就停）、98. Validate Binary Search Tree（中序序列是否嚴格遞增）、426. Convert BST to Sorted Doubly Linked List（中序邊走邊接指標）。

共同前提都是那句話：BST 的中序遍歷 = 遞增序列。整理見 Tree 遍歷模板。

複雜度

預先展開成佇列

• __init__ 時間 $O(n)$（extend 寫法在斜樹會退化成 $O(n^2)$）
• next() / hasNext() 時間 $O(1)$
• 空間 $O(n)$ — 整個中序結果都存著

受控遞迴（棧）

• __init__ 時間 $O(h)$ — 只壓最左邊那條路徑
• next() 時間均攤 $O(1)$ — 單次最壞 $O(h)$，但 n 次呼叫總共只有 $O(n)$
• hasNext() 時間 $O(1)$
• 空間 $O(h)$ — 棧裡最多一條根到葉的路徑

其中 n 是節點數、h 是樹高。兩者的 next() 都是（均攤）$O(1)$，差別完全在空間，而 follow-up 要的就是那個 $O(h)$。