Print tree in a zigzag manner

Question

I want to print a tree in a zigzag manner level-by-level. I'm wondering if there are any logic issues in my code and any further performance improvement ideas by using a smarter algorithm implementation.

class Node:
 def __init__(self, value):
 self.value = value
 self.left = None
 self.right = None
def zigzagPrint(wq, isLeft):
 if not wq:
 return
 buf = []
 if isLeft == True:
 for i in range(len(wq)):
 buf.append(wq[i].value)
 else:
 for i in range(len(wq)-1, -1, -1):
 buf.append(wq[i].value)
 print buf
def zigzag(root):
 if not root:
 return
 wq = []
 wqTemp=[]
 wq.append(root)
 isLeft = True
 while True:
 zigzagPrint(wq, isLeft)
 isLeft = not isLeft
 while len(wq) > 0:
 node = wq.pop(0)
 if node.left:
 wqTemp.append(node.left)
 if node.right:
 wqTemp.append(node.right)
 # inner while
 wq = wqTemp
 wqTemp = []
 if len(wq) == 0:
 return
if __name__ == "__main__":
 n1 = Node(1)
 n21 = Node(2)
 n1.left = n21
 n22 = Node(3)
 n1.right = n22
 n31 = Node(4)
 n32 = Node(5)
 n21.left = n31
 n21.right = n32
 n33 = Node(6)
 n22.right = n33
 n41 = Node(7)
 n42 = Node(8)
 n31.left = n41
 n31.right = n42
 n43 = Node(9)
 n32.right = n43
zigzag(n1)

Expected output

The tree in the code can be represented as

 1
 / \
 2 3
 / \ \
 4 5 6
 / \ \
7 8 9

And the associated output is

[1]
[3, 2]
[4, 5, 6]
[9, 8, 7]

Answer 1

Use better variable names

Right now, I have no idea what you mean by wq. Also, buf could be spelled out as buffer without loosing anything.

Also, PEP8 recommends snake_case for variables and two blank lines between function definitions.

Better comparisons

You should use the fact that isLeft is already a boolean and use if isLeft: instead of if isLeft == True:

Also, you could use that the empty list evaluates as False and save one call to len:

if len(wq) == 0:
 return
if not wq:
 return

Iterate over object, not its indices

Instead of doing:

for i in range(len(wq)):
 buf.append(wq[i].value)
for i in range(len(wq)-1, -1, -1):
 buf.append(wq[i].value)

do:

for node in wq:
 buf.append(node.value)
for node in wq[::-1]:
 buf.append(node.value)

But what you really want is probably:

buf = [x.value for x in wq]
# or
buf = [x.value for x in wq[::-1]]
buf = [x.value for x in reversed(wq)]

where the latter is necessary, because deque does not support the slice notation.

This already saves a lot of append calls and merges them into one assignment. But since these two things are very similar (only the direction changes), we can also write;

buffer = iter(wq) if is_left else reversed(wq)
print [node.value for node in buffer]

Use collections.deque

Since you are always popping left, you could use collections.deque, where this operation is faster than for a simple list.

from collections import deque
...
wq = deque([root])
wq_temp = deque([])

Miscellaneous

For this I think it is better to return as early as possible. Also, use tuple assignment:

# inner while
wq = wqTemp
wqTemp = []
if len(wq) == 0:
 return
# inner while
if not wqTemp:
 return
wq, wqTemp = wqTemp, []

Even better, you can move the check for the empty list to the condition of the while loop, because initially wq is either empty (then the while loop does nothing) or has root in it so is not empty.

This

wq = []
wqtemp = []
wq.append(root)

can be written as:

wq, wqtemp = [root], []

Final result

Wrapping everything together:

from collections import deque
class Node:
 def __init__(self, value):
 self.value = value
 self.left = None
 self.right = None
def zigzag_print(wq, is_left):
 if not wq:
 return
 buffer = iter(wq) if is_left else reversed(wq)
 print [node.value for node in buffer]
def zigzag(root):
 wq, wq_temp = deque([root]), deque([])
 is_left = True
 while wq:
 zigzag_print(wq, is_left)
 is_left = not is_left
 while wq:
 node = wq.popleft()
 if node.left:
 wq_temp.append(node.left)
 if node.right:
 wq_temp.append(node.right)
 wq, wq_temp = wq_temp, deque([])

• Graipher: Hey @LinMa, thanks for the +1. As I said I didn't have a too close look at the algorithm itself, mainly because I couldn't think of a different way at that moment. However, I full-heartedly agree with what MatthiasEttinger suggested in his answer. In the end, you won't find an algorithm for printing a tree which is better than O(n), because you have to print every element. And his algo is O(n) as far as I can see.
• Graipher: @LinMa pattern is whether to go forward through the children (iter just returns an iterator and reversed returns an iterator that starts at the back and goes backwards)
• Graipher: @LinMa Now that you mention it, yeah it should be.
• Graipher: @LinMa Originally I thought because it was implemented differently (faster), but it is not really. You do save some overhead of the for control structure (wiki.python.org/moin/PythonSpeed/PerformanceTips#Loops). You also save the function lookup of list.append. It is also a lot easier to read, if it does not get too long and/or nested.
• Graipher: First, in python 2.x range() will return an actual list, which will take as much memory as the list is long (xrange() returns an iterator that does not have this problem), while in python 3.x range is what xrange is on python 2.x. And there you can already see the advantage of an iterator: it does not care about the underlying structure, as long as it has a next() and StopIteration defined. When the alternative to an iterator is an intermediate list, the memory needed for the latter can be saved.

Answer 2

Building on @Graipher's answer, I don't think you need to explicitly maintain the state of the direction of iteration in the is_left variable. All you need to do is to alternate between iter and reversed to traverse your nodes for printing... and alternate indefinitely.

itertools.cycle is exactly meant for that.

You also don't need to pop (or popleft) items from wq since you use an intermediate list and you override wq with that intermediate list right after iterating over it. An explicit loop will perform better here:

import itertools
class Node:
 def __init__(self, value):
 self.value = value
 self.left = None
 self.right = None
def zigzag(root):
 nodes = [root]
 for pattern in itertools.cycle((iter, reversed)):
 # Print current layer
 if not nodes:
 return
 print [node.value for node in pattern(nodes)]
 # Collect children
 children = []
 for node in nodes:
 if node.left:
 children.append(node.left)
 if node.right:
 children.append(node.right)
 nodes = children

• Lin Ma: Thanks Mathias, appreciate the advice and vote up, not used reverse too much and found it is incredibly useful. Besides your above comments, it seems the algorithm is the same (idea) as mine? I just want to ask if (1) you find any logical errors in my code (2) any performance improvement ideas. Besides your great comments all above. Thanks.
• Lin Ma: Hi Mathias, if you could comment on my above question, it will be great.
• Lin Ma: Hi Mathias, a further question, what ispattern in your code? Thanks.
• 301_Moved_Permanently: @LinMa It's hard to comment further on the logic (beside adding in a list and popping from the other) as I don't know your exact requirements except what you described in your post. If the code is part of a greater project, knowing more about the why and how could lead to further improvements. For instance I base the solution to produce the actual output, but if it is flexible it might be better to use str.join somehow. As for pattern, it is alternatively iter and reversed, meaning going forward, then backward, then forward... when printing the rows.
• Lin Ma: Thanks Mathias, vote up for the reply. It is just a standalone exercise for fun problem and not a bigger project code base. It seems both your algorithm and my algorithm are of O(n) performance? Any specific performance improvement you think your implementation is better than me (just from performance perspective, I know from coding style and other perspective, there are a lot to learn from your code)? Thanks.