Python: Simplify many if-statements

Question

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']

+1 for including a test suite with a refactor request :)

Pavel Anossov
– Pavel Anossov

2013-03-21 00:05:12 +00:00
Commented Mar 21, 2013 at 0:05 — Pavel Anossov
– Pavel Anossov, Commented Mar 21, 2013 at 0:05

Pavel Anossov · Accepted Answer · 2013-03-20 23:58:52Z

5

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.

answered Mar 20, 2013 at 23:58

Pavel Anossov

63.3k16 gold badges156 silver badges125 bronze badges

Sign up to request clarification or add additional context in comments.

1 Comment

abarnert Over a year ago

+1. This is exactly what I meant by my last paragraph. I think it was easier to write in Python than in English. (It's certainly easier to understand that way.)

Bill Lynch · Accepted Answer · 2013-03-20 23:58:51Z

2

The only thing that comes to my mind is this:

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

edited Mar 20, 2013 at 23:58

answered Mar 20, 2013 at 23:53

Bill Lynch

82.4k16 gold badges132 silver badges181 bronze badges

3 Comments

isedev Over a year ago

not really a simplification though... but neater in a way, so +1.

Bill Lynch Over a year ago

I'd argue that it's more clear that each traversal is implemented correctly here than in the original question. But yeah, there's not much that can be done to clarify this code.

abarnert Over a year ago

If you really want: (preorder_traversal, inorder_traversal, postorder_traversal)[order](node) will actually make this "simpler"… but obviously less readable.

abarnert · Accepted Answer · 2013-03-20 23:59:05Z

I don't think there's a good way to eliminate the three if order statements without reorganizing the algorithm. The position at which ret.append happens is dependent on the value, so you pretty much have to check it all three times, one way or another.

But there's one obvious way to reorganize things to remove a couple ifs:

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)

Of course it's one line longer, but it's only 4 ifs instead of 6.

Another possibility is to change things to keep track of all three positions, and insert into the appropriate position after the fact:

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

This removes all of the ifs, but I don't think the result is exactly easier to read or comprehend…

Really, the way to make this easier to read and comprehend is probably to switch to a functional algorithm (recursive mutable algorithms are rarely fun to think through)… but that's just going to make the three if bodies larger, not get rid of them.

John · Accepted Answer · 2013-03-20 23:57:29Z

0

~~Try something like~~

~~if order in (Order.PREORDER, Order.INORDER, ...): ret.append(node.data)~~

Nitpicks

You should have your conditions on multiple lines

if cond:
    pass

instead of

if cond: pass

also consider using is and is not when comparing to None instead of == and !=.

edited Mar 20, 2013 at 23:57

answered Mar 20, 2013 at 23:47

John

13.8k7 gold badges53 silver badges104 bronze badges

3 Comments

isedev Over a year ago

same issue as previous (now deleted) answer: the order of the conditions IS important... what you suggest changes that order.

John Over a year ago

@isedev, +1 you're right, I'll leave it so others don't make the same mistake.

isedev Over a year ago

fair enough, can't remove the -1 (wasn't me).

Collectives™ on Stack Overflow

Python: Simplify many if-statements

4 Answers 4

1 Comment

3 Comments

Comments

3 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

4 Answers 4

1 Comment

3 Comments

Comments

3 Comments

Your Answer

Sign up or log in

Post as a guest

Related