準備過互聯網公司的服務端崗位面試的人，對於二叉樹的三種遍歷方式想必是如數家珍。假設以類BinaryTree定義一棵二叉樹

class BinaryTree:
    def __init__(self, left, right, value):
        self.left = left
        self.right = right
        self.value = value

實現一個前序遍歷的算法便是信手拈來的事情

def preorder_traversal(tree, func):
    """前序遍歷二叉樹的每個節點。"""
    if tree is None:
        return
    func(tree.value)
    preorder_traversal(tree.left, func)
    preorder_traversal(tree.right, func)

隨着行業曲率的增大，要求寫出不使用遞歸的版本也沒什麼過分的

def iterative_preorder_traversal(tree, func):
    nodes = [tree]
    while len(nodes) > 0:
        node = nodes.pop()
        func(node)
        if node.left is not None:
            nodes.append(node.right)
        if node.left is not None:
            nodes.append(node.left)

一直以來，我覺得這種用一個顯式的棧來代替遞歸過程中隱式的棧的做法就是鏡花水月。但最近卻找到了它的一個用武之地——用於實現iterator。

iterator是個啥？

這年頭，iterator已經不是什麼新鮮事物了，許多語言中都有支持，維基百科上有一份清單列出了比較知名的語言的iterator特性。按照Python官方的術語表中的定義，iterator表示一個數據流，反覆調用其__next__方法可以一個接一個地返回流中的下一項數據。將內置函數iter作用於list、str、tuple類型的對象，可以獲得相應的迭代器

$ cat get_iter.py
# -*- coding: utf8 -*-
if __name__ == '__main__':
    values = [
        [1, 2, 3],
        'Hello, world!',
        (True, None),
    ]
    for v in values:
        print('type of iter({}) is {}'.format(v, type(iter(v))))
$ python get_iter.py
type of iter([1, 2, 3]) is <class 'list_iterator'>
type of iter(Hello, world!) is <class 'str_iterator'>
type of iter((True, None)) is <class 'tuple_iterator'>

寫一個前序遍歷的iterator

一個iterator對象必須要實現__iter__和__next__方法：

• __iter__只需要返回iterator對象自身即可；
• 顧名思義，__next__負責返回下一個元素。

仔細觀察一下前文中的iterative_preorder_traversal函數可以看出：

• nodes = [tree]屬於初始化邏輯；
• len(nodes) > 0用於判斷是應當拋出StopIteration，還是應當繼續返回下一個值（nodes.pop()）；
• 最後四行就是負責填充nodes，好讓它可以在下一次調用__next__的時候有值可以返回的。

到這裏，iterator的具體實現代碼已經呼之欲出了

class BinaryTreePreorderIterator:
    def __init__(self, root):
        nodes = []
        if root is not None:
            nodes.append(root)
        self.nodes = nodes

    def __iter__(self):
        return self

    def __next__(self):
        if len(self.nodes) == 0:
            raise StopIteration
        node = self.nodes.pop()
        if node.right is not None:
            self.nodes.append(node.right)
        if node.left is not None:
            self.nodes.append(node.left)
        return node.value

構造一棵這樣的滿二叉樹

用BinaryTreePreorderIterator可以正確地打印出每一個節點的值

if __name__ == '__main__':
    tree = BinaryTree(
        BinaryTree(
            BinaryTree(None, None, 1), 
            BinaryTree(None, None, 3), 
            2,
        ),
        BinaryTree(
            BinaryTree(None, None, 5), 
            BinaryTree(None, None, 7), 
            6,
        ),
        4,
    )
    for n in BinaryTreePreorderIterator(tree):
        print('{}\t'.format(n), end='')
# 打印內容為：4    2    1    3    6    5    7

iterator的優勢

顯然，iterator比起preorder_traversal更為靈活——很容易在for-in循環內添加各種各樣的控制邏輯：用continue跳過一些值，或者用break提前結束遍歷過程。這些在函數preorder_traversal中做起來會比較彆扭。

聰明的你應該已經發現了，大可不必將preorder_traversal拆解到一個構造方法和一個__next__方法中。用generator寫起來明明更加直觀

def preorder_generator(tree):
    """返回一個能夠以前序遍歷的次序遍歷二叉樹節點的generator。"""
    nodes = [tree]
    while len(nodes) > 0:
        node = nodes.pop()
        yield node.value
        if node.left is not None:
            nodes.append(node.right)
        if node.left is not None:
            nodes.append(node.left)

但是，很多語言並不支持generator。與之相比，iterator要親民得多，更容易移植。例如，即使是Common Lisp這種一窮二白的語言，也可以實現和Python的iterator以及for類似的效果

(in-package #:cl-user)

(defpackage #:com.liutos.binary-tree
  (:use #:cl))

(in-package #:com.liutos.binary-tree)

(defclass preorder-iterator ()
  ((nodes
    :initform nil)
   (tree
    :initarg :tree))
  (:documentation "前序遍歷二叉樹的迭代器"))

(defmethod initialize-instance :after ((instance preorder-iterator) &key)
  (with-slots (nodes tree)
      instance
    (when tree
      (push tree nodes))))

(defgeneric next (iterator)
  (:documentation "返回迭代器的下一個值。"))

(define-condition stop-iteration (error)
  ()
  (:documentation "Python中StopIteration異常的等價物。"))

(defmethod next ((iterator preorder-iterator))
  (with-slots (nodes) iterator
    (when (null nodes)
      (error 'stop-iteration))

    (let ((node (pop nodes)))
      ;; 一個節點的結構為：(值 左子樹 右子樹)
      (when (third node)
        (push (third node) nodes))
      (when (second node)
        (push (second node) nodes))
      (first node))))

(defmacro for-in (var iterator &body forms)
  "將iterator中的值逐個綁定到變量var上，並執行forms中的表達式。"
  (let ((iter (gensym)))
    `(let ((,iter ,iterator))
       (handler-case
           (loop
              (let ((,var (next ,iter)))
                ,@forms))
         (stop-iteration (c)
           (declare (ignorable c)))))))

(defparameter *tree*
  '(4 (2 (1 nil nil) (3 nil nil)) (6 (5 nil nil) (7 nil nil))))

(defun test-preorder-iterator ()
  "測試前序遍歷迭代器。"
  (for-in n (make-instance 'preorder-iterator
                           :tree *tree*)
    (format t "~D~C" n #\Tab)))

後記

中序遍歷和後序遍歷也可以寫成迭代器，證明略。

閲讀原文