用 Graphviz 绘制一棵漂亮的二叉树

起因

之前用 Rust 写了一个 AVL 树的实现，就很自然的想把树用可视化的图像画出来，在一波搜索过后，最后都指向了一位叫 Emden Gansner（ Graphviz 主要贡献者之一） 的大佬在 2010 年写的一段脚本，原作者是在邮件列表（链接1，链接2）中回复别人的问题时提供的这段脚本，由于这个邮件列表原来的存档网站已经无法访问，现在能搜到的基本上都是别人对这段脚本引用，比如： stackoverflow。

不过这个 gvpr 我实在是看不懂，所以我希望能够直接使用 dot 来绘制二叉树，这样在生成图像的时候就只需要使用 dot 命令行工具而不需要额外的脚本了。

尝试

但是事情看起来并没有那么简单，假如我们现在有一个文件 tree.dot：

digraph G {
    node [shape=circle]
    edge [arrowhead=vee]
    8 -> 4
    4 -> 2
    2 -> 1
    2 -> 3
    4 -> 6
    6 -> 5
    6 -> 7
    8 -> 10
    10 -> 9
    10 -> 12
    12 -> 11
}

直接用 dot tree.dot -Tsvg -otree_0.svg 生成出来的图像是这样的：

然后我们可以过一遍那位大佬的脚本（把脚本保存为 tree.g）：

dot tree.dot | gvpr -c -ftree.g | neato -n -Tsvg -otree_1.svg

得到的图像如下：

看起来效果还行，但是 11 这个节点原本应该在左边，却放到了右边，虽然我可以修改 dot 文件让节点放到左边去，不过我想要的效果是不依赖这个脚本。

原理

定位左右子节点

在左右子节点中间添加一个虚拟的占位节点，并且使之与父节点在同一竖直方向上。这可以通过为占位节点（node）和边（edge）设置 style=invis 属性，并且给父节点和占位节点设置同一个 group 来实现。

同时为了更好的控制左右节点之间的间距，可以对全图设置 nodesep=0，然后对占位节点设置 label="", width=0.3，通过修改这里的 0.3 来控制左右节点的间距。当然也可以选择为占位节点设置 width=0，然后通过修改全图设置 nodesep 来控制左右节点的间距，分别使用两种方式设置相同的节点间距所需的数值有一个大约 2 倍多的比例关系。

所以一开始的 tree.dot 更新为：

digraph G {
    graph [nodesep=0.1]
    node [shape=circle]
    edge [arrowhead=vee]
    8 [group=8]
    4 [group=4]
    8 -> 4
    2 [group=2]
    4 -> 2
    2 -> 1
    _2 [group=2, label="", width=0, style=invis]
    2 -> _2 [style=invis]
    2 -> 3
    _4 [group=4, label="", width=0, style=invis]
    4 -> _4 [style=invis]
    6 [group=6]
    4 -> 6
    6 -> 5
    _6 [group=6, label="", width=0, style=invis]
    6 -> _6 [style=invis]
    6 -> 7
    _8 [group=8, label="", width=0, style=invis]
    8 -> _8 [style=invis]
    10 [group=10]
    8 -> 10
    10 -> 9
    _10 [group=10, label="", width=0, style=invis]
    10 -> _10 [style=invis]
    12 [group=12]
    10 -> 12
    12 -> 11
    _12 [group=12, label="", width=0, style=invis]
    12 -> _12 [style=invis]
}

所得图像为：

如果我们使所有的占位节点及边可见，得到的图像是：

看起来效果也还行，不过仔细看还是会发现一点问题，4 这个节点的位置偏右，这棵数节点不够多还不是很明显，但是如果是一棵大树，这个偏移会十分明显以及丑陋。

定位父节点

为了解决前面的问题，我们可以考虑将 4 所对应的占位节点下移到 3 和 5 中间，不过 10 所对应的占位节点就不需要下移。

整个规则总结起来就是：一个节点对应的占位节点应该与该节点的左子树的最大节点和右子树的最小节点中距离较近的那一个处于同一层。如果根据这个规则，占位节点位于紧邻的下一层，我们可以不用额外设置。

所以最后的 tree.dot 文件应该是这样的（注意含 rank=same 的两行）：

digraph G {
    graph [nodesep=0.1]
    node [shape=circle]
    edge [arrowhead=vee]
    8 [group=8]
    4 [group=4]
    8 -> 4
    2 [group=2]
    4 -> 2
    2 -> 1
    _2 [group=2, label="", width=0, style=invis]
    2 -> _2 [style=invis]
    2 -> 3
    _4 [group=4, label="", width=0, style=invis]
    4 -> _4 [style=invis]
    6 [group=6]
    4 -> 6
    6 -> 5
    _6 [group=6, label="", width=0, style=invis]
    6 -> _6 [style=invis]
    6 -> 7
    {rank=same; _4; 5}
    _8 [group=8, label="", width=0, style=invis]
    8 -> _8 [style=invis]
    10 [group=10]
    8 -> 10
    10 -> 9
    _10 [group=10, label="", width=0, style=invis]
    10 -> _10 [style=invis]
    12 [group=12]
    10 -> 12
    12 -> 11
    _12 [group=12, label="", width=0, style=invis]
    12 -> _12 [style=invis]
    {rank=same; _8; 9}
}

最终得到的图像是：

代码实现

我在之前写的 avl_tree 中实现了将树直接导出为这样的 dot 文件的函数：

pub fn print_dot<T: PartialOrd + Display>(tree: &AvlTreeNode<T>) {
    fn print_node<T: PartialOrd + Display>(node: &TreeNode<T>) {
        let mut target = None;
        let mut distance = 0;

        if let Some(x) = &node.left {
            let mut left_max = x;
            let mut left_distance = 1;
            while let Some(x) = &left_max.right {
                left_max = x;
                left_distance += 1;
            }
            target = Some(&left_max.val);
            distance = left_distance;

            if x.left.is_some() || x.right.is_some() {
                println!("    {} [group={}]", x.val, x.val);
            }
            println!("    {} -> {}", node.val, x.val);
            print_node(x);
        }

        if node.left.is_some() || node.right.is_some() {
            println!(
                "    _{} [group={}, label=\"\", width=0, style=invis]",
                node.val, node.val
            );
            println!("    {} -> _{} [style=invis]", node.val, node.val);
        }

        if let Some(x) = &node.right {
            let mut right_min = x;
            let mut right_distance = 1;
            while let Some(x) = &right_min.left {
                right_min = x;
                right_distance += 1;
            }
            if right_distance <= distance {
                target = Some(&right_min.val);
                distance = right_distance;
            }

            if x.left.is_some() || x.right.is_some() {
                println!("    {} [group={}]", x.val, x.val);
            }
            println!("    {} -> {}", node.val, x.val);
            print_node(x);
        }

        if distance > 1 {
            if let Some(x) = target {
                println!("    {{rank=same; _{}; {}}}", node.val, x);
            }
        }
    }

    if let Some(x) = tree {
        println!("digraph G {{");
        println!("    graph [nodesep=0.1]");
        println!("    node [shape=circle]");
        println!("    edge [arrowhead=vee]");
        if x.left.is_some() || x.right.is_some() {
            println!("    {} [group={}]", x.val, x.val);
        }
        print_node(x);
        println!("}}");
    }
}

最后附上一棵随机生成的 50 个节点的 AVL 树：

完