Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Python: Simplify many if-statements

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']
like image 796
Mr. Polywhirl Avatar asked Apr 28 '26 00:04

Mr. Polywhirl


1 Answers

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.

like image 164
Pavel Anossov Avatar answered Apr 30 '26 14:04

Pavel Anossov



Donate For Us

If you love us? You can donate to us via Paypal or buy me a coffee so we can maintain and grow! Thank you!