非递归遍历二叉树到底有什么用

准备过互联网公司的服务端岗位面试的人，对于二叉树的三种遍历方式想必是如数家珍。假设以类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)))

后记

中序遍历和后序遍历也可以写成迭代器，证明略。

阅读原文