joowani / binarytree

Python Library for Studying Binary Trees
http://binarytree.readthedocs.io
MIT License
1.81k stars 171 forks source link

Feature: Export to graphviz dot format #19

Closed villmow closed 3 years ago

villmow commented 5 years ago

First of all nice package, thanks for that! I just wanted to visualize a binary tree with Graphviz thats why I adapted this utility function from the treelib package, maybe it fits somewhere in your library. As I wanted to have string values for nodes and that doesn't work with your package I supplied an optional id2name dictionary, which stores the mapping.

import binarytree as bt

def binarytree_to_dot(root: bt.Node, filename: str, id2name: Optional[Dict[int, str]] = None, 
                      shape: str = 'ellipse', graph: str = 'digraph') -> None:
    """Exports the tree in dot format of the graphviz software. 
    Left childs are denoted by a full stroke arrow, right childs by a dashed arrow."""

    if not isinstance(root, bt.Node):
        raise ValueError("Root must be a binarytree.")

    if id2name is None:
        id2name = {node.value: str(node.value) for node in root}

    nodes, connections = [], []

    for node in root.levelorder:
        nid = node.value
        state = '"' + str(nid) + '"' + ' [label="' + str(id2name[nid]) + '", shape=' + shape + ']'
        nodes.append(state)

        if node.left is not None:
            cid = node.left.value
            connections.append('"' + str(nid) + '":sw' + ' -> ' + '"' + str(cid) + '"')
        if node.right is not None:
            cid = node.right.value
            connections.append('"' + str(nid) + '":se' + ' -> ' + '"' + str(cid) + '"  [style=dashed]')

    # write nodes and connections to dot format
    with open(filename, 'wt') as f:
        f.write(graph + ' tree {\n')
        for n in nodes:
            f.write('\t' + n + '\n')

        f.write('\n')
        for c in connections:
            f.write('\t' + c + '\n')

        f.write('}')

Example Use

Use it like this

root = bt.tree(height=6)
binarytree_to_dot(root, filename="test_tree.dot")

After installing graphviz, you can plot the tree with the following command:

dot -T png -o test_tree.png test_tree.dot

which will result in the following image: image

joowani commented 5 years ago

Hi @villmow,

Thanks for liking binarytree and this is pretty cool work!

I have received a few requests for graphviz integration in the past. Right now, however, because of my schedule, I've paused all feature work on my open source projects (PR submissions will also be very slow). If you want to share your work with more people I encourage you to publish your own library built on top of binarytree. I will keep this open for any future updates.

joowani commented 3 years ago

Instead of implementing this method from scratch, I've decided to use python-graphviz instead (it has many more features) in version 6.0.0. Please see the documentation for more details!