我有一个函数,它遍历一棵树,并将元素作为列表返回。是否有简化所有if语句的方法在treeToList::traverse中,因为它看起来有点冗余?
#!/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']