CODEWITHSUNDEEP
pythonrecursiontree
KOB
KOB

Reputation: 4545

Algorithm to produce a clean __str__ method for a tree

I have a Tree class which is made up of Node objects. Each Node has the following instance vairables:

The Node class also has a method get_ancenstors() which returns a list of the Nodes in the path from this Node's parent to the root.

The tree is built using a recursive function and has the following instance variables:

The Tree class has a method dfs() which performs a depth first search/traversal of the tree and yields every Node visited. It also has a __str__() method as such:

def __str__(self):
    string = ''
    for dir in self.dfs(self.root):
        # TODO: Create branches as '|-' 
        branches = ('  ' * (dir.depth))
        string += f'{branches}{dir.name}\n'
    return string

Which produces an ouput like:

Dir 1
  Dir 2
    Dir 3
  Dir 4
    Dir 5
      Dir 6
      Dir 7
    Dir 8
  Dir 9
    Dir 10

Instead, I want a visually cleaner, but more complex implementation of this __str__() method which produces an output like this:

Dir 1
|--- Dir 2
|    |--- Dir 3
|
|--- Dir 4
|    |--- Dir 5
|    |    |--- Dir 6
|    |    |--- Dir 7
|    |
|    |--- Dir 8
|
|--- Dir 9
     |--- Dir 10
     |--- Dir 11
          |--- Dir 12

Can anyone help with the logic behind this?

Upvotes: 0

Views: 117

Answers (2)

Ajax1234
Ajax1234

Reputation: 71451

A possibility is to create a recursive generator that yields a complete row for every child. Then, in __str__, you can simply use str.join:

class Tree:
   ...
   def to_dir(self, d = 0, flag=True):
      if flag:
        yield self.name if not d else '|'+"|".join(["\t"]*(d-1))+f'{"|"*(bool(d-1))}--- {self.name}'
      else:
         yield "".join(["\t"]*(d-1))+f'{"|"*(not d)}{"|"*(bool(d))}--- {self.name}'
      for i, a in enumerate(self.children):
         yield from a.to_dir(d=d+1, flag = falg)
   def __str__(self):
      return '\n'.join(self.to_dir())

Example:

class Tree:
   def __init__(self, n, c):
      self.name, self.children = n, c
   def to_dir(self, d = 0, flag=True): 
      if flag:
         yield self.name if not d else '|'+"|".join(["\t"]*(d-1))+f'{"|"*(bool(d-1))}--- {self.name}'
      else:
         yield "".join(["\t"]*(d-1))+f'{"|"*(not d)}{"|"*(bool(d))}--- {self.name}'
      for i, a in enumerate(self.children):
         yield from a.to_dir(d=d+1, flag=False if not flag else i != len(self.children)-1)
   def __str__(self):
      return '\n'.join(self.to_dir())

t = Tree('Dir 1', [Tree('Dir 2', [Tree('Dir 5', [])]), Tree('Dir 3', []), Tree('Dir 4', [Tree('Dir 6', [Tree('Dir 7', [])])])])

Output:

Dir 1
|--- Dir 2
       |--- Dir 5
|--- Dir 3
|--- Dir 4
       |--- Dir 6
              |--- Dir 7

Upvotes: -1

RomanPerekhrest
RomanPerekhrest

Reputation: 92854

Use the following solution with specific formatting logic:

def __str__(self):
    string = ''
    for dir in self.dfs(self.root):
        branches, depth = '', dir.depth
        if depth > 0:
            branches = '|--- '  # initial (closest) dashed block
            branches = '|    ' * (depth - 1) + branches
        string += f'{branches}{dir.name}\n'
    return string

Simplified test case for demo:

from collections import namedtuple

Node = namedtuple('Node', 'name depth')
dfs = [Node(f'Dir{i+1}', depth=i) for i in range(4)]

def n__str__(dfs):
    string = ''
    for dir in dfs:
        branches, depth = '', dir.depth
        if depth > 0:
            branches = '|--- '  # initial (closest) dashed block
            branches = '|    ' * (depth - 1) + branches
        string += f'{branches}{dir.name}\n'
    return string

print(n__str__(dfs))

The output:

Dir1
|--- Dir2
|    |--- Dir3
|    |    |--- Dir4

Upvotes: 0

Related Questions