Python：简化多个if语句

Question

我有一个函数，它可以遍历一棵树，并把树里的元素以列表的形式返回。请问有没有办法简化一下treeToList::traverse里面的所有if语句？因为看起来有点重复。

#!/usr/bin/python

def enum(**enums):
  return type('Enum', (), enums)

Order = enum(PREORDER=0, INORDER=1, POSTORDER=2)
def treeToList(root, order=Order.INORDER):
  ret = list()
  def traverse(node, order):
    if order == Order.PREORDER: ret.append(node.data)
    if node.right != None: traverse(node.right, order)
    if order == Order.INORDER: ret.append(node.data)
    if node.down != None: traverse(node.down, order)
    if order == Order.POSTORDER: ret.append(node.data)
  traverse(root, order)
  return ret

class node:
  def __init__(self, data=None):
    self.data = data
    self.down = None
    self.right = None

if __name__ == '__main__':
  root = node('F')
  root.right = node('B')
  root.down = node('G')
  root.right.right = node('A')
  root.right.down = node('D')
  root.down.down = node('I')
  root.right.down.right = node('C')
  root.right.down.down = node('E')
  root.down.down.right = node('H')

  print treeToList(root, Order.PREORDER)
  print treeToList(root, Order.INORDER)
  print treeToList(root, Order.POSTORDER)

输出

['F', 'B', 'A', 'D', 'C', 'E', 'G', 'I', 'H']
['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I']
['A', 'C', 'E', 'D', 'B', 'H', 'I', 'G', 'F']

Answer 1

我脑海中想到的只有这个：

def treeToList(root, order=Order.INORDER):
    ret = list()
    def inorder_traversal(node):
        if node is not None:
            inorder_traversal(node.right)
            ret.append(node.data)
            inorder_traversal(node.down)

    def preorder_traversal(node):
        if node is not None:
            ret.append(node.data)
            preorder_traversal(node.right)
            preorder_traversal(node.down)

    def postorder_traversal(node):
        if node is not None:
            postorder_traversal(node.right)
            postorder_traversal(node.down)
            ret.append(node.data)

    if order == Order.PREORDER:
        preorder_traversal(node)
    elif order == Order.INORDER:
        inorder_traversal(node)
    else:
        postorder_traversal(node)

    return ret

Answer 2

我觉得要想去掉这三个 if order 语句，没什么好办法，除非重新组织一下算法。因为 ret.append 的位置是根据值来决定的，所以你基本上得检查三次，无论怎么做。

不过，有一种明显的方法可以重新安排一下，去掉几个 if：

def traverse(node, order):
    if node is None: return
    if order == Order.PREORDER: ret.append(node.data)
    traverse(node.right, order)
    if order == Order.INORDER: ret.append(node.data)
    traverse(node.down, order)
    if order == Order.POSTORDER: ret.append(node.data)

当然，这样会多一行代码，但 if 的数量从 6 个减少到 4 个。

还有一种可能是改变方式，跟踪所有三个位置，然后在最后插入到合适的位置：

def traverse(node, order):
    if node is None: return
    prepos = len(ret)
    traverse(node.right, order)
    inpos = len(ret)
    traverse(node.down, order)
    postpos = len(ret)
    pos = (prepos, inpos, postpos)[order]
    ret[pos:pos+1] = node.data

这样可以去掉所有的 if，但我觉得结果并没有变得更容易阅读或理解……

其实，要让这个更容易理解，可能需要换成一种函数式算法（递归的可变算法通常不太好理解）……不过这样做只是会让三个 if 的内容变得更大，并不能去掉它们。

Answer 3

好吧，如果你去掉闭包的话，一个纯函数可能会更清楚：

def treeToList(node, order=Order.INORDER):
    if node is None:
        return []

    right = treeToList(node.right, order)
    down = treeToList(node.down, order)
    current = [node.data]

    if order == Order.PREORDER:
        return current + right + down

    if order == Order.INORDER:
        return right + current + down

    if order == Order.POSTORDER:
        return right + down + current

不过，这样会产生很多中间列表。