1
\$\begingroup\$

Here is my code to print tree in a level by level way. It needs to print an empty node as well to keep tree structure like full binary tree for elegancy reasons.

I'm wondering if anything you think could improve, including the efficiency and not sure if any bugs (I did testing, not found bugs). I'm especially wondering if for part of code -- "print the remaining all empty leaf node part" could be optimized (I mean if any smarted way to optimize the while loop to cover all logic, and skip this code snippet especially handling the last level).

The tree I tested in this example looks like:

 1
/ \
2  3
/\ /\
4  5 6

Here is the related output:

[1]
[2, 3]
[4, '$', 5, 6]
['$', '$', '$', '$', '$', '$', '$', '$']

class Node:
    def __init__(self):
        self.value = 0
        self.left = None
        self.right = None

def printTree(root):
    buf = []
    output = []
    if not root:
        print '$'
    else:
        buf.append(root)
        count = 1
        nextCount = 0
        while count > 0:
            node = buf.pop(0)
            if node:
                output.append(node.value)
                count -= 1
            else:
                output.append('$')
            if node and node.left:
                buf.append(node.left)
                nextCount += 1
            else:
                buf.append(None)
            if node and node.right:
                buf.append(node.right)
                nextCount += 1
            else:
                buf.append(None)
            if count == 0:
                print output
                output = []
                count = nextCount
                nextCount = 0
        # print the remaining all empty leaf node part
        for i in range(len(buf)):
            output.append('$')
        print output

if __name__ == "__main__":
    root = Node()
    root.value = 1
    leftNode = Node()
    leftNode.value = 2
    root.left = leftNode
    rightNode = Node()
    rightNode.value = 3
    root.right = rightNode
    leftleftNode = Node()
    leftleftNode.value = 4
    leftNode.left = leftleftNode
    rightleftNode = Node()
    rightleftNode.value = 5
    rightNode.left = rightleftNode
    rightrightNode = Node()
    rightrightNode.value = 6
    rightNode.right = rightrightNode

    printTree(newRoot)
\$\endgroup\$
1
\$\begingroup\$

A few small changes: Python has a styleguide, PEP8, which programmers are recommended to follow. One of its recommendations is to use lower_case names for both variables and functions. camelCase is also accepted, if applied consistently (which it is in your code).

Your if conditions in the while loop repeatedly check if node exists. Just put all of them in if node:. Additionally, the checks for node.left and node.right are basically identical, they can be combined by looping over for n in [node.left, node.right]:.

Whenever you do a buffer.pop(0) you probably want a collections.deque. While for lists, popping on the left is O(n) and O(1) on the right, both operations are O(1) for a deque.

Python has the nice feature of multiple assignments, allowing things like:

count, nextCount = nextCount, 0

You are always creating an instance of a node and directly afterwards you assign the value. It is better to add a parameter for the value to the class definition:

class Node:
    def __init__(self, value=0):
        self.value = value
        self.left = None
        self.right = None

With these changes I get this code:

from collections import deque

class Node:
    def __init__(self, value=0):
        self.value = value
        self.left = None
        self.right = None


def printTree(root):
    buf = deque()
    output = []
    if not root:
        print '$'
    else:
        buf.append(root)
        count, nextCount = 1, 0
        while count:
            node = buf.popleft()
            if node:
                output.append(node.value)
                count -= 1
                for n in (node.left, node.right):
                    if n:
                        buf.append(n)
                        nextCount += 1
                    else:
                        buf.append(None)
            else:
                output.append('$')
            if not count:
                print output
                output = []
                count, nextCount = nextCount, 0
        # print the remaining all empty leaf node part
        output.extend(['$']*len(buf))
        print output

if __name__ == "__main__":
    root = Node(1)
    leftNode = Node(2)
    root.left = leftNode
    rightNode = Node(3)
    root.right = rightNode
    leftleftNode = Node(4)
    leftNode.left = leftleftNode
    rightleftNode = Node(5)
    rightNode.left = rightleftNode
    rightrightNode = Node(6)
    rightNode.right = rightrightNode

    printTree(newRoot)

This is still not very efficient code. If you give Node a magic method to decide if it is True or False, your code can be greatly simplified. This method is called __nonzero__ in python 2.x and __bool__ in python3.x.

class Node:
    def __init__(self, value='$', left=None, right=None):
        self.value = value
        self.left = left
        self.right = right

    def __nonzero__(self):
        return self.value != '$'

    def __str__(self):
        buf, out = [self], []
        while buf:
            out.append("{}".format([node.value for node in buf]))
            if any(node for node in buf):
                children = []
                for node in buf:
                    for subnode in (node.left, node.right):
                        children.append(subnode if subnode else Node())
                buf = children
            else:
                break
        return "\n".join(out)

if __name__ == "__main__":
    root = Node(1, Node(2, Node(4)), Node(3, Node(5), Node(6)))
    print root

This gets rid of all the counting and so on. It does, however, not print the last line of all missing nodes (which seemed odd to me anyways). Normally I would use None for an empty Node, not '$', but it also works this way.

I also added left and right, as parameters to the Node, allowing a nested definition, which might be less readable if it becomes too big, but might come handy, when quickly generating a tree like here.

If you also add the magic method __str__ to Node, you can just do print root (or any other node and all its subnodes).

If you insist on having each level printed like a list (to ease copy&pasting it), replace the line

out.append(" ".join(str(node.value) for node in buf))

with

out.append("{}".format([node.value for node in buf]))

If you want to learn more about magic methods, have a look here.

| improve this answer | |
\$\endgroup\$
  • \$\begingroup\$ Thanks Graipher, vote up. 1. I see you also have code output.extend(['$']*len(buf)) outside while loop, is there a way to put it inside the loop in your mind (I might be wrong but I think one loop code might be more elegant)? 2. Any logical bugs or issues in your mind of my original code? \$\endgroup\$ – Lin Ma Aug 10 '16 at 15:14
  • 1
    \$\begingroup\$ @LinMa See updated answer, at the end. I propose a slightly different implementation (the algorithm is similar). \$\endgroup\$ – Graipher Aug 12 '16 at 9:34
  • 1
    \$\begingroup\$ I still don't see the sense in printing this row, because it will by definition always contain only empty nodes. however, have a look at the edited code. \$\endgroup\$ – Graipher Aug 15 '16 at 8:24
  • 1
    \$\begingroup\$ BTW, Graipher, re-studied your code today. I think in each iteration of loop, you call any I think any is an O(n) operation in performance perspective, not sure in this way your algorithm will become overall performance O(n*n)? Thanks. \$\endgroup\$ – Lin Ma Aug 17 '16 at 8:02
  • 1
    \$\begingroup\$ @Linma any is O(n) at worst, i.e. if none of the conditions is true. Otherwise it is O(m), where m is the position of the first true condition, m <= n. So yes, it might be slightly slower for a tree like Node(0, None, Node(1, None, Node(2, None, Node(3, None, ...)))), i.e. where the left node is always empty. \$\endgroup\$ – Graipher Aug 17 '16 at 11:04

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.