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.
45 lines
1.6 KiB
Python
45 lines
1.6 KiB
Python
3 years ago
|
from graph_utils import make_name_generator
|
||
|
|
from graphviz import Digraph, Graph
|
||
|
|
|
||
|
|
def make_node (graph, generate_node_name, label):
|
||
|
|
node_name = generate_node_name()
|
||
|
|
graph.node(node_name, label=label)
|
||
|
|
return node_name
|
||
|
|
|
||
|
|
|
||
|
|
def visualize_node (graph, make_node, generate_node_name, node, previous_node_name = None, tailport = None):
|
||
|
|
if node:
|
||
|
|
node_name = make_node(graph, generate_node_name, node.value)
|
||
|
|
|
||
|
|
if previous_node_name:
|
||
|
|
graph.edge(previous_node_name, node_name, tailport=tailport, headport='s', penwidth='1')
|
||
|
|
|
||
|
|
|
||
|
|
visualize_node(graph, make_node, generate_node_name, node.left, node_name, 'nw')
|
||
|
|
# Hack to have a more beautiful layout
|
||
|
|
visualize_node(graph, make_node, generate_node_name, None, node_name, 'n')
|
||
|
|
visualize_node(graph, make_node, generate_node_name, node.right, node_name, 'ne')
|
||
|
|
|
||
|
|
else:
|
||
|
|
# Invisible node to balance the tree.
|
||
|
|
node_name = generate_node_name()
|
||
|
|
graph.node(node_name, label=node_name, style='invis')
|
||
|
|
graph.edge(previous_node_name, node_name, tailport=tailport, headport='s', style='invis')
|
||
|
|
|
||
|
|
def make_graph (graphname):
|
||
|
|
graph = Graph(name=graphname, format='svg', engine='dot')
|
||
|
|
graph.attr('graph', splines='line', rankdir='BT')
|
||
|
|
return graph
|
||
|
|
|
||
|
|
def visualize (tree, graphname, make_node=make_node, generate_node_name = make_name_generator(length=3)):
|
||
|
|
graph = make_graph(graphname)
|
||
|
|
visualize_node(graph, make_node, 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')
|