Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Print tree with indentations

I want to print out a simple tree that describes a band. I start by creating a node called "Band" and then I make the child "Wind instruments" which in turns have the children "Saxophone" and "Trumpet". Then I make a sibling to "Wind instruments" called "Song" and so on. The code is pretty straightforward:

class Node:                                     
    value = ""
    down = None
    right = None
def write(p):
    if p==None:
        return
    print(p.value)
    if p.down!=None:        #My idea is that if we go down the tree, we indent first
        print("    ",end="")     
        write(p.down)
    write(p.right)      #If we don't go down, we simply write out the siblings
a=Node()
a.value="Band"
a.down=Node()
a.down.value="Wind instruments"
a.down.down=Node()
a.down.down.value="Saxophone"
a.down.down.right=Node()
a.down.down.right.value="Trumpet"
a.down.right=Node()
a.down.right.value="Song"
a.down.right.right=Node()
a.down.right.right.value="String instruments"
a.down.right.right.down=Node()
a.down.right.right.down.value="Guitar"
a.down.right.right.down.right=Node()
a.down.right.right.down.right.value="Bass"
write(a)

The output is:

Band
    Wind instruments
    Saxophone
Trumpet
Song
String instruments
    Guitar
Bass

But I want the output to be:

Band
    Wind instruments
        Saxophone
        Trumpet
    Song
    String instruments
        Guitar
        Bass

Anyone have an idea how to achieve this?

like image 816
Lozansky Avatar asked Jan 06 '23 11:01

Lozansky


1 Answers

To make an indented print depending on the recursion level, the trick is to use an argument keeping the level you're at when recursing:

# default with a level of 0, and an indent of 4 characters
def write(p, depth=0, indent=4):
    if p==None:
        return
    # here we multiply the level by the number of indents
    # and then you multiply that number with a space character
    # which will magically show as that number of spaces.
    print("{}{}".format(" "*(indent*depth), p.value))
    if p.down!=None:
        # then you do not need your print(…, end='') hack
        # simply increase the depth
        write(p.down, depth=depth+1, indent=indent)
    # and for siblings, do not increase the depth
    write(p.right, depth=depth, indent=indent)

The trick I'm using here is that per default the level is 0, and as you're going deeper, you're increasing the depth by passing the argument increased by 1.

Then, when you want to printout the indentation, all you have to do is multiply the indentation string with that value (and the indent size) so then you get to print the indent as you like:

>>> "A"*42
'AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA'

So as a result:

>>> write(a)
Band
    Wind instruments
        Saxophone
        Trumpet
    Song
    String instruments
        Guitar
        Bass

And if you want to make it narrower, because you're having a lot of recursion:

>>> write(a, indent=1)
Band
 Wind instruments
  Saxophone
  Trumpet
 Song
 String instruments
  Guitar
  Bass

as a bonus, I'd advice you to make your write() function a method of your Node class. And if you rename it __str__:

class Node:
    value = ""
    down = None
    right = None

    # converts a node into a string
    def as_str(self, depth=0, indent=4):
        # building the current node's line, and add it to a list
        ret = ["{}{}".format(" "*(indent*depth), self.value)]
        if self.down:
            # append down recursion first to the list
            ret.append(self.down.as_str(depth=depth+1, indent=indent))
        if self.right:
            # then append right recursion to the list
            ret.append(self.right.as_str(depth=depth, indent=indent))
        # build the new string, joining each element of the list with a newline
        return "\n".join(ret)

    # a handler for printing the list nicely
    def __str__(self):
        return self.as_str()

    def as_repr(self, depth=0, max_depth=2):
        # building the current node's line, and add it to a list
        ret = ["'{}'".format(self.value)]
        if depth > max_depth:
            ret.append("…")
        else:
            if self.down:
                # append down recursion first to the list
                ret.append(self.down.as_repr(depth=depth+1, max_depth=max_depth))
            if self.right:
                # then append right recursion to the list
                ret.append(self.right.as_repr(depth=depth, max_depth=max_depth))
            # build the new string, joining each element of the list with a newline
        return "Node<{}>".format(",".join(ret))

    # you might want to also make the repr() nicer
    def __repr__(self):
        return self.as_repr()

And as a result:

>>> a
Node<'Band',Node<'Wind instruments',Node<'Saxophone',…>,Node<'Song',Node<'String instruments',Node<'Guitar',…>>>>>
>>> print(a)
Band
    Wind instruments
        Saxophone
        Trumpet
    Song
    String instruments
        Guitar
        Bass

HTH

like image 101
zmo Avatar answered Jan 15 '23 22:01

zmo