二叉树python实现和可视化

在刷LeetCode时，有时候需要在本地调试代码，但是苦于本地没有树的数据类型，所以自己动手用python写了一个二叉树类，并且实现了可视化。

下面的程序仅仅是为了创建二叉树，方便在刷LeetCode有关题目时进行本地调试代码，所以有些功能没有加上去，比如删除节点的功能。程序虽然是二叉树类，也容易扩展到多叉树。

目前程序只能实现无重复节点的二叉树构建。

我看到网上有不少的用python实现二叉树的代码，但是它们构建出来的几乎都是特定类型的二叉树，比如二叉搜索树、完全二叉树等。本文的程序可以创建任意类型的二叉树。

程序：

import collections
import networkx as nx
import matplotlib.pyplot as plt


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


class Tree:
    def __init__(self, root_val):
        root_node = Node(root_val)
        self.root = root_node
    
    def add(self, val, parent_val, position):
        parent = self.find(parent_val)   # find the parent node. We assume that there is no duplicate nodes.
        node = Node(val)
        if position == 0:
            parent.left = node
        if position == 1:
            parent.right = node

    def find(self, value):
        que = collections.deque([self.root])
        while que:
            node = que.popleft()
            if node.val == value:
                return node
            if node.left:
                que.append(node.left)
            if node.right:
                que.append(node.right)
        raise KeyError('value not found')


    # draw the tree. https://zhuanlan.zhihu.com/p/35574577
    def draw(self): 
        graph = nx.DiGraph()
        graph, pos = self.create_graph(graph, self.root)
        fig, ax = plt.subplots(figsize=(8, 10))  # 比例可以根据树的深度适当调节
        nx.draw_networkx(graph, pos, ax=ax, node_size=300)
        plt.show()

    def create_graph(self, G, node, pos={}, x=0, y=0, layer=1):
        pos[node.val] = (x, y)
        if node.left:
            G.add_edge(node.val, node.left.val)
            l_x, l_y = x - 1 / 2 ** layer, y - 1
            l_layer = layer + 1
            self.create_graph(G, node.left, x=l_x, y=l_y, pos=pos, layer=l_layer)
        if node.right:
            G.add_edge(node.val, node.right.val)
            r_x, r_y = x + 1 / 2 ** layer, y - 1
            r_layer = layer + 1
            self.create_graph(G, node.right, x=r_x, y=r_y, pos=pos, layer=r_layer)
        return (G, pos)


# 示例：
tree = Tree(6)
tree.add(val=2, parent_val=6, position=0)  # position: 0 means 'left', 1 means 'right'
tree.add(3, 6, 1)
tree.add(13, 2, 0)
tree.add(4, 2, 1)
tree.add(12, 4, 0)
tree.add(0, 4, 1)
tree.add(24, 3, 0)
tree.add(17, 3, 1)

tree.draw()

程序运行结果：