HiveBrain v1.2.0
Get Started
← Back to all entries
patternpythonMinor

Print tree in a zigzag manner

Submitted by: @import:stackexchange-codereview··
0
Viewed 0 times
printzigzagmannertree

Problem

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]

Solution

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([])

Code Snippets

if len(wq) == 0:
    return

if not wq:
    return
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)
for node in wq:
    buf.append(node.value)
for node in wq[::-1]:
    buf.append(node.value)
buf  = [x.value for x in wq]
# or
buf = [x.value for x in wq[::-1]]
buf = [x.value for x in reversed(wq)]
buffer = iter(wq) if is_left else reversed(wq)
print [node.value for node in buffer]

Context

StackExchange Code Review Q#135245, answer score: 4

Revisions (0)

No revisions yet.