用pyplot打印二叉树，实现二叉树的可视化

第一次写博客哈哈哈哈哈
最近实验题目多涉及二叉树，每次调试在监视窗口查看二叉树需要挨个挨个点，十分麻烦，于是就想直接将二叉树画出来看看，比较直观。

思路

为了避免重合，二叉树左右子树的距离是关键
让每一个节点占用一个列空间，这样就不会重合了
所以
左节点与父节点在X轴上的距离为 左节点 的 右子树宽度+1 乘以一个水平距离常数
右节点与父节点在X轴上的距离为 右节点 的 左子树宽度+1 乘以一个水平距离常数

每当画好一个节点点，确定其左右孩子在X轴上的距离，再画这个节点连接孩子的边

首先要有一棵二叉树

random_tree.py

二叉树定义

class TreeNode():
    def __init__(self, x):
        self.val = x
        self.left = None
        self.right = None

根据中序遍历生成随机的二叉树

代码很简单，我就不废话了

def generate_random_tree_by_mid(mid):
    '''根据中序序列生成随机BT'''
    if len(mid) == 0:
        return None
    root = TreeNode(random.choice(mid))
    index = mid.index(root.val)
    root.left = generate_random_tree_by_mid(mid[:index])
    root.right = generate_random_tree_by_mid(mid[index+1:])
    return root

生成二叉树或者二叉搜索树的中序序列

BST的中序是一个有序的序列，若不生成BST，则将列表打乱

def generate_random_tree(n, isBST = False):
    '''生成随机形状和随机值的BT或BST'''
    numbers = [x for x in range(1, 200)]
    mid = []
    for _ in range(n):
        mid.append(random.choice(numbers))
    if isBST:
        mid.sort()
    else:
        random.shuffle(mid)

    root = generate_random_tree_by_mid(mid)
    return root

定义获取画二叉树必要信息的方法

关于树的方法，基本上可以用递归实现

get_params.py

from random_tree import TreeNode

def get_left_width(root):
    '''获得根左边宽度'''
    return get_width(root.left)

def get_right_width(root):
    '''获得根右边宽度'''
    return get_width(root.right)

def get_width(root):
    '''获得树的宽度'''
    if root == None:
        return 0
    return get_width(root.left) + 1 + get_width(root.right)

def get_height(root):
    '''获得二叉树的高度'''
    if root == None:
        return 0
    return max(get_height(root.left), get_height(root.right)) + 1

最后来画图

这里用到了matplotlib.pyplot来画图
圆形来自 matplotlib.patches中的Circle

show_BTree.py

基本参数设置

d_hor = 4   #节点水平距离
d_vec = 8   #节点垂直距离
radius = 2  #节点的半径

获取树的高度和宽度

调用get_params.py的方法获取树的高度和宽度，用来设置画图的尺寸
尺寸跟随树的 宽 高 动态变化

def get_w_h(root):
    '''返回树的宽度和高度'''
    w = get_width(root)
    h = get_height(root)
    return w, h

画节点和边的方法

画点需要点的x, y 坐标， 点的值

def draw_a_node(x, y, val, ax):
    '''画一个节点'''
    c_node = Circle((x,y), radius=radius, color='green')
    ax.add_patch(c_node)
    plt.text(x, y, '%d' % val, ha='center', va= 'bottom',fontsize=11)

def draw_a_edge(x1, y1, x2, y2, r=radius):
    '''画一条边'''
    x = (x1, x2)
    y = (y1, y2)
    plt.plot(x, y, 'k-')

初始化画布大小，返回第一个节点的坐标

def create_win(root):
    '''创建窗口'''
    WEIGHT, HEIGHT = get_w_h(root)
    WEIGHT = (WEIGHT+1)*d_hor
    HEIGHT = (HEIGHT+1)*d_vec
    fig = plt.figure(figsize=(11, 9))
    ax = fig.add_subplot(111)
    plt.xlim(0, WEIGHT)
    plt.ylim(0, HEIGHT)

    x = (get_left_width(root) + 1) * d_hor #x, y 是第一个要绘制的节点坐标，由其左子树宽度决定
    y = HEIGHT - d_vec
    return fig, ax, x, y     

最关键的代码

采用中序遍历，根据其他节点与父节点的相对位置，画出节点

def print_tree_by_inorder(root, x, y, ax):
    '''通过中序遍历打印二叉树'''
    if root == None:
        return
    draw_a_node(x, y, root.val, ax)
    lx = rx = 0
    ly = ry = y - d_vec
    if root.left != None:
        lx = x - d_hor * (get_right_width(root.left) + 1)   #x-左子树的右边宽度
        draw_a_edge(x, y, lx, ly, radius)
    if root.right != None:
        rx = x + d_hor * (get_left_width(root.right) + 1)   #x-右子树的左边宽度
        draw_a_edge(x, y, rx, ry, radius)
    #递归打印    
    print_tree_by_inorder(root.left, lx, ly, ax)
    print_tree_by_inorder(root.right, rx, ry, ax)

主函数中调用即可

先画窗口，再画树

def show_BTree(root):
    '''可视化二叉树'''
    _, ax, x, y = create_win(root)
    print_tree_by_inorder(root, x, y, ax)
    plt.show()
    
def main():
    root = random_tree.generate_random_tree(50, True)
    show_BTree(root)


if __name__ == '__main__':
    main()

样例打印了一棵有50个节点的BST

源码

https://download.youkuaiyun.com/download/steven_l1994/10837482