5
\$\begingroup\$

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]
\$\endgroup\$
4
+50
\$\begingroup\$

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([])
\$\endgroup\$
  • 1
    \$\begingroup\$ 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. \$\endgroup\$ – Graipher Jul 24 '16 at 14:46
  • 1
    \$\begingroup\$ @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) \$\endgroup\$ – Graipher Jul 24 '16 at 18:09
  • 1
    \$\begingroup\$ @LinMa Now that you mention it, yeah it should be. \$\endgroup\$ – Graipher Jul 24 '16 at 18:19
  • 1
    \$\begingroup\$ @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. \$\endgroup\$ – Graipher Jul 25 '16 at 8:15
  • 1
    \$\begingroup\$ 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. \$\endgroup\$ – Graipher Jul 25 '16 at 10:15
4
\$\begingroup\$

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
\$\endgroup\$
  • \$\begingroup\$ 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. \$\endgroup\$ – Lin Ma Jul 19 '16 at 18:06
  • \$\begingroup\$ Hi Mathias, if you could comment on my above question, it will be great. \$\endgroup\$ – Lin Ma Jul 21 '16 at 23:36
  • \$\begingroup\$ Hi Mathias, a further question, what ispattern in your code? Thanks. \$\endgroup\$ – Lin Ma Jul 24 '16 at 17:56
  • 1
    \$\begingroup\$ @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. \$\endgroup\$ – Mathias Ettinger Jul 24 '16 at 19:05
  • \$\begingroup\$ 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. \$\endgroup\$ – Lin Ma Jul 24 '16 at 22:31

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.