dsa.pretty_print

  1""" Module to access functions for a clearer visual representation of certain data structures. """
  2import math
  3
  4def heap_print(heap):
  5    """
  6    Print a heap from root to leaves.
  7
  8    Args:
  9        heap: The heap object to print.
 10    """
 11    if not heap or heap.count() == 0:
 12        return
 13
 14    _array_print(heap.count(), heap.enumerate())
 15
 16
 17def _array_print(node_count, enum_sequence):
 18    """
 19    Print a heap from root to leaves.
 20
 21    Args:
 22        heap: The heap object to print.
 23    """
 24    height = math.floor(math.log2(node_count))
 25    level_str = ""
 26    current_level = 0
 27    value_width = 3
 28    max_width = 2 ** (height - 1) * value_width
 29
 30    for index, node in enum_sequence:
 31        level = int(math.log2(index + 1))
 32        columns = 2 ** (level - 1)
 33        column_width = int(max_width / columns)
 34        if current_level != level:
 35            current_level = level
 36            print(level_str)
 37            level_str = ""
 38        level_str += f"{node:^{column_width}}"
 39    print(level_str)
 40    print()
 41
 42
 43def tree_to_array(node, index: int=0, tree_array=None):
 44    """
 45    (helper function) Create an array filled with index and value pairs from a node based tree.
 46
 47    Args:
 48        node: The starting node.
 49        index (int): The starting index.
 50        tree_array: The destination array.
 51
 52    Returns:
 53        Array filled with tree values.
 54    """
 55    if not tree_array:
 56        tree_array = []
 57    if node is None:
 58        return
 59    tree_array.append((index, node.value))
 60    tree_to_array(node.left, index * 2 + 1, tree_array)
 61    tree_to_array(node.right, index * 2 + 2, tree_array)
 62    
 63    return tree_array
 64
 65def get_tree_height(node) -> int:
 66    """
 67    (helper function) Calculate the height of a tree.
 68
 69    Args:
 70        node: The starting node.
 71
 72    Returns:
 73        The height of a tree.
 74    """
 75    if node is None:
 76        return 0
 77    else:
 78        return max(get_tree_height(node.left) + 1, get_tree_height(node.right) + 1)
 79
 80def fill_complete_tree(tree):
 81    """
 82    (helper function) Force a binary tree to be a complete tree by filling any empty nodes.
 83
 84    Args:
 85        tree: The tree to fill.
 86
 87    Returns:
 88        A new tree that is complete.
 89    """
 90    if tree is None:
 91        return None
 92    
 93    tree_array = tree_to_array(tree.root)
 94    if tree_array is None:
 95        return
 96    
 97    tree_height = get_tree_height(tree.root)
 98
 99    # build empty complete tree
100    array_size = (2 ** tree_height) - 1
101    new_tree = [ "" ] * array_size
102
103    # fill the complete tree array
104    for index, value in tree_array:
105        new_tree[index] = value
106    return new_tree
107
108def tree_print(tree):
109    """
110    Print a tree from root to leaves.
111
112    Args:
113        tree: The tree object to print.
114
115    Notes:
116        Reuses heap_print() by converting tree into a complete tree array.
117    """
118    complete_tree = fill_complete_tree(tree)
119    if complete_tree:
120        _array_print(len(complete_tree), enumerate(complete_tree))