利用 graphviz 输出红黑树

发布于 Feb. 17, 2017, 6:06 p.m.

0 个评论

因为当时没有好好学算法，所以现在开始补习算法。看到红黑树的时候，觉得可视化红黑树能够帮助我们更直观地了解其结构，于是在网上搜索之后，发现可以使用 graphviz 配合python 接口 pygraphviz 来实现。

安装

需要注意，使用 brew 安装 graphviz 之后，需要对 pip 指定参数来安装 pygraphvix。

1 2	brew install graphviz pip install --install-option="--include-path=/usr/local/include/" --install-option="--library-path=/usr/local/lib/" pygraphviz

代码

红黑树代码略。

# coding=utf-8
import pygraphviz as pgv
import itertools

class Node:
    def __init__(self, key=None, left=None, right=None, parent=None, color=RED):
        self.key = key
        self.left = left
        self.right = right
        self.parent = parent
        self.color = color

NIL = Node(color=BLACK)

class Tree:
    def __init__(self):
        self.nil = NIL
        self.root = NIL

    # other methods are omitted ...

    def draw(self, filename='tree.png', draw_nil=0):
        if self.root == self.nil:
            return
        graph = pgv.AGraph(strict=False, directed=True)
        graph.node_attr.update({
            'shape':     'circle',
            'style':     'filled',
            'fontcolor': 'white',
        })
        graph.add_node(self.root.key, color=self.root.color)
        if draw_nil == 0:
            self.draw_subtree(self.root, graph)
        elif draw_nil == 1:
            # use one nil node for leaves and parent of root
            graph.add_node('nil1', color='black', label='nil')
            graph.add_edge('nil1', self.root.key)
            self.draw_subtree(self.root, graph, draw_nil)
        else:
            # use different nil nodes for leaves and parent of root
            graph.add_node('nil1', color='black', label='nil')
            graph.add_edge('nil1', self.root.key)
            self.draw_subtree(self.root, graph, draw_nil, counter=itertools.count(2))
        # print(graph.string())
        graph.draw(filename, prog='dot')

    def draw_subtree(self, node, graph, draw_nil=0, counter=None):
        if node == self.nil:
            return
        if node.left != self.nil:
            graph.add_node(node.left.key, color=node.left.color)
            graph.add_edge(node.key, node.left.key, tailport='sw')
            self.draw_subtree(node.left, graph, draw_nil, counter)
        elif draw_nil == 1:
            graph.add_edge(node.key, 'nil1', tailport='sw')
        elif draw_nil == 2 and counter is not None:
            nil_id = 'nil%d' % next(counter)
            graph.add_node(nil_id, color='black', label='nil')
            graph.add_edge(node.key, nil_id, tailport='sw')

        if node.right != self.nil:
            graph.add_node(node.right.key, color=node.right.color)
            graph.add_edge(node.key, node.right.key, tailport='se')
            self.draw_subtree(node.right, graph, draw_nil, counter)
        elif draw_nil == 1:
            graph.add_edge(node.key, 'nil1', tailport='se')
        elif draw_nil == 2 and counter is not None:
            nil_id = 'nil%d' % next(counter)
            graph.add_node(nil_id, color='black', label='nil')
            graph.add_edge(node.key, nil_id, tailport='se')


if __name__ == '__main__':
    tree = Tree()
    values = [41, 38, 31, 12, 19, 8]
    for x in values:
        tree.rb_insert(Node(x))
    tree.draw(filename='tree_nil_0.png', draw_nil=0)
    tree.draw(filename='tree_nil_1.png', draw_nil=1)
    tree.draw(filename='tree_nil_2.png', draw_nil=2)

draw 方法的 draw_nil 参数默认为 0（不画 nil 结点），设置为 1 则叶子和根父亲共用一个 nil 结点，其他值则叶子和根父亲有各自的 nil 结点。下图分别为上述代码执行后的 tree_nil_0.png，tree_nil_1.png 和 tree_nil_2.png。

不画 nil 结点的红黑树