I have a function that traverses a tree and returns the elements as a list. Is there a way to simplify all the if statements in treeToList::traverse, because it looks sort of redundant?
#!/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)
Output
['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']
Well, if you get rid of the closure.. a pure function is probably clearer:
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
but builds a lot of intermediate lists of course.
If you love us? You can donate to us via Paypal or buy me a coffee so we can maintain and grow! Thank you!
Donate Us With