solution_237.py

#!/usr/bin/python3
# --------------------------------------------------------------------
# display a tree structure
# (build a test tree and then display it)
# --------------------------------------------------------------------
# Base in: stackoverflow.com/questions/20242479/
#          printing-a-tree-data-structure-in-python
# --------------------------------------------------------------------
# links
# a. stackoverflow.com/questions/34012886/
#    print-binary-tree-level-by-level-in-python
# b. pypi.org/project/PrettyPrintTree/
# --------------------------------------------------------------------

import random

# --------------------------------------------------------------------
# ---- class definition: tree node
# --------------------------------------------------------------------

class Node:

    def __init__(self,parent,sym):
        self.parent = parent
        self.sym    = sym
        self.left   = None
        self.right  = None

    def __str__(self):
        return str(self.sym)

# --------------------------------------------------------------------
# ---- add a node to a sub tree below a parent node
# --------------------------------------------------------------------

def add_node(node,parent):

    if node.sym < parent.sym:
        if parent.left == None:
            parent.left = node
            node.parent = parent
            return
        add_node(node,parent.left)
        return
    
    if parent.right == None:
        parent.right = node
        node.parent  = parent
        return
    add_node(node,parent.right)
    return        

# --------------------------------------------------------------------
# ---- traverse a tree - results is a list of tree nodes
# --------------------------------------------------------------------

def traverse_tree(node,nodes):

    if node is None: return
    traverse_tree(node.left,nodes)
    if nodes is not None: nodes.append(node)
    traverse_tree(node.right,nodes)
    return
        
# --------------------------------------------------------------------
# ---- create a balanced tree
# ---- built a tree from a list of nodes ordered by value
# ---- (set up and then call a recursive function)
# ---- based on:
# ---- www.geeksforgeeks.org/convert-normal-bst-balanced-bst/
# --------------------------------------------------------------------

def construct_balanced_tree(root):

    def balanced_tree_util(nodes,start,end,parent):

        if start > end: return None
        
        # ---- get middle node

        mid = (start + end)//2        
        node = nodes[mid]

        # ---- construct/add to balanced tree
        # ---- a. fill in root node's value
        # ---- b. add a new node to the tree

        if parent.sym is None:
            parent.sym = node.sym
            parent.parent = None
        else:
            add_node(Node(None,node.sym),parent)            
        
        # ---- construct left and right subtrees

        balanced_tree_util(nodes,start,mid-1,parent)
        balanced_tree_util(nodes,mid+1,end,parent)

        return

    # ---- get a list of original tree nodes

    nodes = []
    traverse_tree(root,nodes)

    # ---- create a sorted list of node from
    # ---- the original tree node list
    
    sorted_nodes = sorted(nodes, key=lambda node:node.sym)
    
    # ---- create a new balanced tree

    new_root = Node(None,None)

    n = len(sorted_nodes)
    balanced_tree_util(sorted_nodes,0,n-1,new_root)

    # ---- return the balanced tree

    return new_root

# --------------------------------------------------------------------
# ---- print node
# --------------------------------------------------------------------

def print_node(node,title=None):

    if title is None: title = 'Node: '
    
    if node is None: print(f'{title}None'); return
    
    print(f'{title}sym={node.sym} ' +\
          f'parent={node.parent} '  +\
          f'left={node.left} '      +\
          f'right={node.right}')

    return

# --------------------------------------------------------------------
# ---- print a "new fashion" fancy tree
# --------------------------------------------------------------------

def print_fancy_tree(node):

    ##elbow = '+--'
    ##pipe  = '|  '
    ##tee   = '|--'
    ##blank = '   '
    elbow = '└──'
    pipe  = '│  '
    tee   = '├──'
    blank = '   '

    def print_fancy_tree_util(node,header):

        if node.left is not None:
            if node.right is not None:
                print(header+tee+node.left.sym)
                headerl = header + pipe
            else:
                print(header+elbow+node.left.sym)
                headerl = header + blank
            print_fancy_tree_util(node.left,headerl)

        if node.right is not None:
            if node.left is not None:
                print(header+elbow+node.right.sym)
                headerr = header + blank
            else:
                print(header+elbow+node.right.sym)
                headerr = header + blank
            print_fancy_tree_util(node.right,headerr)

        return

    if node is None: return

    print(node.sym)

    print_fancy_tree_util(node,'')

    return

# --------------------------------------------------------------------
# ---- print an "old fashion" dash tree
# --------------------------------------------------------------------

def print_dash_tree(node, dash='-', level=0):

    if node is None: return

    print(f'[{level}] ' + dash*(level*3) +  '(' + str(node.sym) + ')')

    level += 1
    print_dash_tree(node.left,dash,level)
    print_dash_tree(node.right,dash,level)

    return

# --------------------------------------------------------------------
# ---- build random test tree of symbols
# --------------------------------------------------------------------

def build_random_test_tree(symbols:list):

    if len(symbols) < 1: return None

    root = Node(None,symbols[0])

    if len(symbols) > 1:
        build_random_sub_trees(root,symbols,0)

    return root

# --------------------------------------------------------------------
# ---- build sub trees under a parent node
# ----
# ---- Note:
# ---- 1. idx is the symbols list index of the most recent
# ----    symbol used - the next symbol is idx+1
# ---- 2. random functions are there to mix it up a little
# --------------------------------------------------------------------

def build_random_sub_trees(parent:Node,
                           symbols:list,
                           idx:int) -> int:

    # ---- safety "stuff" (basically don't do anything)

    if parent is None: return idx         # no parent node?

    if idx+1 >= len(symbols): return idx  # any symbol remaining?

    # ---- randomly select which, if any, parent nodes are None
    # ---- left  node is None? # right node is None

    nleft = nright = False
    if random.randint(0,100) < 50: nleft  = True
    if random.randint(0,100) < 50: nright = True 

    # ---- create parent right node node only

    if nleft and not nright:           

        if idx+1 < len(symbols):
            idx += 1
            node = Node(parent,symbols[idx])
            parent.right = node
            ##print(f'creating node {symbols[idx]} (idx={idx} right)')

    # ---- create left parent node node only

    elif nright and not nleft:

        if idx+1 < len(symbols):
            idx += 1
            node = Node(parent,symbols[idx])
            parent.left = node
            ##print(f'creating node {symbols[idx]} (idx={idx} left)')

    # ---- create parent left and right nodes

    else:

        if idx+1 < len(symbols):
            idx += 1
            node = Node(parent,symbols[idx])
            parent.left = node
            ##print(f'creating node {symbols[idx]} (idx={idx} left)')

        if idx+1 < len(symbols):
            idx += 1
            node = Node(parent,symbols[idx])
            parent.right = node
            ##print(f'creating node {symbols[idx]} (idx={idx} right)')

    # ---- build sub trees under the parent's left and right nodes
    # ---- randomly switch the order of building the sub trees

    if random.randint(0,100) < 50:
        if parent.left is not None:
            idx = build_random_sub_trees(parent.left,symbols,idx)
        if parent.right is not None:            
            idx = build_random_sub_trees(parent.right,symbols,idx)
    else:
        if parent.right is not None:
            idx = build_random_sub_trees(parent.right,symbols,idx)
        if parent.left is not None:
            idx = build_random_sub_trees(parent.left,symbols,idx)

    return idx     # last symbol to be used

# --------------------------------------------------------------------
# ---- main
# --------------------------------------------------------------------

if __name__ == '__main__':
    
    symbols  = [ 'a','b','c','d','e','f','g','h','i','j' ]

    ##random.seed(20011944)

    # ---- create a random test tree --------------------------------

    print()
    root = build_random_test_tree(symbols)

    print()
    print('-------------- Nodes -------------------')
    print_node(root)
    print_node(root.left)
    print_node(root.right)

    # ---- print "old fashion" dash tree

    print()
    print('-------------- Dash Tree ---------------')
    print_dash_tree(root)

    # ---- print "new fashion" fancy tree

    print()
    print('------------ Fancy Dash Tree -----------')
    print_fancy_tree(root)

    # ---- create a balanced tree from the random tree

    broot = construct_balanced_tree(root)
    print()
    print('---------- Dash Balanced Tree ----------')
    print_dash_tree(broot)

    print()
    print('---------- Fancy Balanced Tree ---------')
    print_fancy_tree(broot)