How to print a tree in Python?

Question

I have the following class which represents a node of a tree:

class Node:
    def __init__(self, name, parent=None):
        self.name = name
        self.parent = parent
        self.children = []
        # ...

        if parent:
            self.parent.children.append(self)

How to print such a tree?

Senega

Answer 1

This is my solution:

def print_tree(current_node, indent="", last='updown'):

    nb_children = lambda node: sum(nb_children(child) for child in node.children) + 1
    size_branch = {child: nb_children(child) for child in current_node.children}

    """ Creation of balanced lists for "up" branch and "down" branch. """
    up = sorted(current_node.children, key=lambda node: nb_children(node))
    down = []
    while up and sum(size_branch[node] for node in down) < sum(size_branch[node] for node in up):
        down.append(up.pop())

    """ Printing of "up" branch. """
    for child in up:     
        next_last = 'up' if up.index(child) is 0 else ''
        next_indent = '{0}{1}{2}'.format(indent, ' ' if 'up' in last else '│', " " * len(current_node.name))
        print_tree(child, indent=next_indent, last=next_last)

    """ Printing of current node. """
    if last == 'up': start_shape = '┌'
    elif last == 'down': start_shape = '└'
    elif last == 'updown': start_shape = ' '
    else: start_shape = '├'

    if up: end_shape = '┤'
    elif down: end_shape = '┐'
    else: end_shape = ''

    print '{0}{1}{2}{3}'.format(indent, start_shape, current_node.name, end_shape)

    """ Printing of "down" branch. """
    for child in down:
        next_last = 'down' if down.index(child) is len(down) - 1 else ''
        next_indent = '{0}{1}{2}'.format(indent, ' ' if 'down' in last else '│', " " * len(current_node.name))
        print_tree(child, indent=next_indent, last=next_last)

Example of use:

shame = Node("shame")

conscience = Node("conscience", shame)
selfdisgust = Node("selfdisgust", shame)
embarrassment = Node("embarrassment", shame)

selfconsciousness = Node("selfconsciousness", embarrassment)
shamefacedness = Node("shamefacedness", embarrassment)
chagrin = Node("chagrin", embarrassment)
discomfiture = Node("discomfiture", embarrassment)
abashment = Node("abashment", embarrassment)
confusion = Node("confusion", embarrassment)

print_tree(shame)

And this is the output:

     ┌conscience
     ├self-disgust
shame┤
     │             ┌self-consciousness
     │             ├shamefacedness
     │             ├chagrin
     └embarrassment┤
                   ├discomfiture
                   ├abashment
                   └confusion

UPDATE:

I pushed a more complete solution on PyPi.

Senega

Great! You found it by yourself. – Tine

Your tree needs a self-esteem boost :) – Dianthus

I just share my morning work as StackOverflow encourages it :) – Senega

Nice solution! as a note to anyone using it on python 2, you'll need # -*- coding: utf-8 -*- at the top of the file – Hundredfold

@Hundredfold you can just use # coding: utf-8 it's easier to remember and it works ;-) – Simeon

Do you have a binary tree adaption? – Haifa

@Haifa you can add only two children per node? – Senega

@Haifa oh sorry, no it won't, the order is not preserved because the algorithm tries to balance the tree. But if the nodes are initialized in the right order, you just have to replace the first six lines by something like that: up, down = [node.children[0]], [node.children[1]] and it will be good. – Senega