You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
56 lines
1.8 KiB
Python
56 lines
1.8 KiB
Python
from string import ascii_uppercase
|
|
from graphviz import Digraph, Graph
|
|
|
|
# Returns a function to generate node names with an optional prefix
|
|
def make_name_generator (length = 1, prefix = ''):
|
|
wheels = [{ 'position': None, 'max': len(ascii_uppercase), 'values': list(ascii_uppercase)} for _ in range(length)]
|
|
|
|
def name_generator ():
|
|
for wheel in wheels:
|
|
if wheel['position'] is None:
|
|
wheel['position'] = 0
|
|
else:
|
|
wheel['position'] += 1
|
|
if wheel['position'] < wheel['max']:
|
|
break
|
|
else:
|
|
wheel['position'] = 0
|
|
|
|
return prefix + ''.join(reversed([wheel['values'][wheel['position']] for wheel in wheels]))
|
|
|
|
return name_generator
|
|
|
|
def visualize (tree, graphname):
|
|
generate_node_name = make_name_generator(length=3)
|
|
|
|
def visualize_node (graph, node, previous_node_name = None, tailport = None):
|
|
if node:
|
|
node_name = generate_node_name()
|
|
graph.node(node_name, label=str(node.value))
|
|
|
|
if previous_node_name:
|
|
graph.edge(previous_node_name, node_name, tailport=tailport, headport='s')
|
|
|
|
|
|
visualize_node(graph, node.left, node_name, 'nw')
|
|
# Hack to have a more beautiful layout
|
|
visualize_node(graph, None, node_name, 'n')
|
|
visualize_node(graph, node.right, node_name, 'ne')
|
|
|
|
else:
|
|
node_name = generate_node_name()
|
|
graph.node(node_name, label=node_name, style='invis')
|
|
graph.edge(previous_node_name, node_name, style='invis', tailport=tailport, headport='s')
|
|
|
|
graph = Graph(name=graphname, format='svg', engine='dot')
|
|
graph.attr('graph', splines='line', rankdir='BT')
|
|
visualize_node(graph, tree)
|
|
graph.render(graphname)
|
|
|
|
if __name__ == '__main__':
|
|
from treesort import make_tree
|
|
from random import randint
|
|
|
|
values = [randint(0, 250) for _ in range(15)]
|
|
tree = make_tree(values)
|
|
visualize(tree, 'tree') |