Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Displaying a tree in ASCII

Tags:

python

tree

As a time-pass activity, I decided to implement a Tree(like) structure in python.
I implemented a Node class (which alone serves the purpose here) like so:

class Node:
    def __init__(self, name, parent, *data):
        self.name = name
        self.parent = parent
        self.data = data
        self.children = []
        self.is_root = False

    def __repr__(self):
        return 'Node '+repr(self.name)

    def dic(self):
        retval = {self:[]}
        for i in self.children:
            retval[self].append(i.dic())
        return retval

    def display(self): # Here
        pass

    def has_children(self):
        return bool(self.children)

    def get_parent(self):
        return self.parent

    def add_child(self, name, *data):
        child = Node(name, self,*data)
        self.children.append(child)
        return child

As you can see the display function is not implemented.
Here's an example tree.

A = Node('A',Node)
A.is_root = True
B = A.add_child('B')
D = B.add_child('D')
C = A.add_child('C')
E = C.add_child('E')
F = C.add_child('F')
G = C.add_child('G')

Here's some sample output for display.

>>> A.display()
    A
  +-^-+
  B   C
  | +-+-+
  D E F G
>>> C.display()
   C
 +-+-+
 E F G

In the shortest form,
How can I "build" an ASCII tree (like above) from the Node class??

In a longer form,
The "Logic" of printing is:

  1. When there is only one child, | is put above the child. (D)
  2. Else, Every child has a + above it, (B,C,E,F)
  3. When there are even no. of children, ^ is put below the parent. (A)
  4. Else, (there are odd no. of children) + is put below the parent. (C)

I have been thinking of starting from below. I realized that there has to be a call to the each of the children, but have been unable to implement anything (of that sorts or otherwise) that gave anything close to it.

like image 535
pradyunsg Avatar asked Mar 28 '13 06:03

pradyunsg


2 Answers

Here's a solution that covers most of what you're looking for.

Like any tree algorithm, recurse down the children of the tree, and combine results at each node. Here's the trick: display() returns a rectangle of text, for example:

aaaaaa
aaaaaa
aaaaaa

Most of the rectangle will be whitespace. Returning only rectangles of text makes it easy to combine results. We'll use the following two helper functions, one to measure block widths, and the other to combine blocks horizontally into larger blocks:

def block_width(block):
    try:
        return block.index('\n')
    except ValueError:
        return len(block)

def stack_str_blocks(blocks):
    """Takes a list of multiline strings, and stacks them horizontally.

    For example, given 'aaa\naaa' and 'bbbb\nbbbb', it returns
    'aaa bbbb\naaa bbbb'.  As in:

    'aaa  +  'bbbb   =  'aaa bbbb
     aaa'     bbbb'      aaa bbbb'

    Each block must be rectangular (all lines are the same length), but blocks
    can be different sizes.
    """
    builder = []
    block_lens = [block_width(bl) for bl in blocks]
    split_blocks = [bl.split('\n') for bl in blocks]

    for line_list in itertools.izip_longest(*split_blocks, fillvalue=None):
        for i, line in enumerate(line_list):
            if line is None:
                builder.append(' ' * block_lens[i])
            else:
                builder.append(line)
            if i != len(line_list) - 1:
                builder.append(' ')  # Padding
        builder.append('\n')

    return ''.join(builder[:-1])

See where this is going? Children return a rectangle that displays themselves and their descendants, and each node will combine these rectangles into a larger rectangle that contains itself. The rest of the code just renders the dashes and pluses:

class Node:
    def display(self): # Here
        if not self.children:
            return self.name

        child_strs = [child.display() for child in self.children]
        child_widths = [block_width(s) for s in child_strs]

        # How wide is this block?
        display_width = max(len(self.name),
                    sum(child_widths) + len(child_widths) - 1)

        # Determines midpoints of child blocks
        child_midpoints = []
        child_end = 0
        for width in child_widths:
            child_midpoints.append(child_end + (width // 2))
            child_end += width + 1

        # Builds up the brace, using the child midpoints
        brace_builder = []
        for i in xrange(display_width):
            if i < child_midpoints[0] or i > child_midpoints[-1]:
                brace_builder.append(' ')
            elif i in child_midpoints:
                brace_builder.append('+')
            else:
                brace_builder.append('-')
        brace = ''.join(brace_builder)

        name_str = '{:^{}}'.format(self.name, display_width)
        below = stack_str_blocks(child_strs)

        return name_str + '\n' + brace + '\n' + below

    # SNIP (your other methods)

And we're off to the races!

                             a                             
+-+-+---------------------------+                          
b e f                           g                          
+     +-+-------------------------+                        
c     h i                         k                        
+       + +-+-+-+-------------+-------------+-+------+     
d       j l m p r             s             O P      Q     
            + +   +-+-+-+---------+             +-----+    
            n q   t u w x         y             R     S    
            +       +     +-------+-------+       +---+---+
            o       v     z       A       M       T   U   Z
                            +-+-+-+-+-+-+ +           +   +
                            B D E H I K L N           V   a
                            +   +   +               +-+-+ +
                            C   F   J               W X Y b
                                +                          
                                G                          

(Requirements like "placing a ^ below the parent" are left as an exercise for the reader)

like image 182
Stu Gla Avatar answered Oct 09 '22 21:10

Stu Gla


I'd like to suggest to take a look at ETE module http://ete.cgenomics.org which implements the functionality you describe here and much more.

At the same time, I'd like to provide this entry Pretty Print output in a sideways tree format in console window where a similar question has been asked before. As you can see in such discussion, the _asciiArt function from ETE provides what, I think, you are looking for.

I hope this helps,

like image 35
scapella Avatar answered Oct 09 '22 20:10

scapella