Python实现普通二叉树

Python实现普通二叉树

二叉树是每个节点最多有两个子树的树结构，本文使用Python来实现普通的二叉树。

关于二叉树的介绍，可以参考：https://blog.csdn.net/weixin_43790276/article/details/104737870

一、实现节点类

所有树结构都是由一个一个的节点构成的，本文使用链式的方式来实现二叉树，所以先实现一个节点类。

# coding=utf-8
class Node(object):
    """节点类"""
    def __init__(self, data, left_child=None, right_child=None):
        self.data = data
        self.parent = None
        self.left_child = left_child
        self.right_child = right_child

在二叉树中添加节点时，要先创建节点，有了节点类，实例化一个节点类的实例即可，节点初始化时是一个孤立的节点，指向的父节点、左子节点、右子节点都为空。将节点“挂”到树上(添加到树中)后，才属于树的一部分。

二、实现二叉树类

实现一个二叉树的类 BinaryTree ，创建二叉树时，实例化一个 BinaryTree 类的实例即可。

class BinaryTree(object):
    """二叉树类"""
    def __init__(self):
        self.__root = None
        self.prefix_branch = '├'
        self.prefix_trunk = '|'
        self.prefix_leaf = '└'
        self.prefix_empty = ''
        self.prefix_left = '─L─'
        self.prefix_right = '─R─'

    def is_empty(self):
        return not self.__root

    @property
    def root(self):
        return self.__root

    @root.setter
    def root(self, data):
        self.__root = Node(data)

    def show_tree(self):
        if self.is_empty():
            print('空二叉树')
            return
        print('-' * 20)
        print(self.__root.data)
        self.__print_tree(self.__root)
        print('-' * 20)

    def __print_tree(self, node, prefix=None):
        if prefix is None:
            prefix, prefix_left_child = '', ''
        else:
            prefix = prefix.replace(self.prefix_branch, self.prefix_trunk)
            prefix = prefix.replace(self.prefix_leaf, self.prefix_empty)
            prefix_left_child = prefix.replace(self.prefix_leaf, self.prefix_empty)
        if self.has_child(node):
            if node.right_child is not None:
                print(prefix + self.prefix_branch + self.prefix_right + str(node.right_child.data))
                if self.has_child(node.right_child):
                    self.__print_tree(node.right_child, prefix + self.prefix_branch + ' ')
            else:
                print(prefix + self.prefix_branch + self.prefix_right)
            if node.left_child is not None:
                print(prefix + self.prefix_leaf + self.prefix_left + str(node.left_child.data))
                if self.has_child(node.left_child):
                    prefix_left_child += '  '
                    self.__print_tree(node.left_child, self.prefix_leaf + prefix_left_child)
            else:
                print(prefix + self.prefix_leaf + self.prefix_left)

    def has_child(self, node):
        return node.left_child is not None or node.right_child is not None

    def __str__(self):
        return str(self.__class__)

一棵二叉树，只要先找到树的根节点(Root)，就可以依次根据节点的父子关系找到所有节点，所以初始化一棵二叉树时，只要先指向树的根节点即可，在没有添加节点到二叉树中时，树是一棵空二叉树，树的根指向为空。

在初始化二叉树时，树的根设置为私有属性，这样做是有必要的，保证根节点的稳定。但是又需要通过实例对象来操作树的根，所以使用 @property 加 @root.setter 提供一对方法供实例对象使用，可以获取或设置二叉树的根节点。

同时，还实现了判断二叉树是否为空的 is_empty() 方法和按树形结构打印二叉树的 show_tree() 方法，在实现完全二叉树的另一篇文章里面有说明。可以参考：https://blog.csdn.net/weixin_43790276/article/details/104033490

三、二叉树添加节点

完全二叉树添加节点是从上到下、从左到右添加的，普通二叉树添加节点没有规律，节点的添加是按需添加的，二叉树的子节点分为左子节点和右子节点，所以实现两个添加节点的方法，分别添加左子节点和右子节点。

    def add_left_child(self, parent, value):
        """添加左子节点"""
        if parent is None:
            print("父节点不存在")
            return
        if parent.left_child is not None:
            print("父节点已有左子节点")
            return
        node = value if isinstance(value, Node) else Node(value)
        parent.left_child = node
        node.parent = parent

    def add_right_child(self, parent, value):
        """添加右子节点"""
        if parent is None:
            print("父节点不存在")
            return
        if parent.right_child is not None:
            print("父节点已有右子节点")
            return
        node = value if isinstance(value, Node) else Node(value)
        parent.right_child = node
        node.parent = parent

add_left_child(parent, value)：添加左子节点。先给定一个父节点，如果父节点已经有了左子节点，则不能添加，如果父节点的左子节点不存在，则将新节点添加到左子节点的位置。将父节点的 left_child 指向新节点，将新节点的 parent 指向父节点。这个方法支持传入一个新节点，也支持传入新节点中保存的数据。

add_right_child(parent, value)：添加右子节点。原理同添加左子节点。

if __name__ == '__main__':
    tree = BinaryTree()
    tree.root = 'A'
    node1 = Node('B')
    tree.add_left_child(tree.root, node1)
    node2 = Node('C')
    tree.add_right_child(tree.root, node2)
    node3 = Node('D')
    node4 = Node('E')
    node5 = Node('F')
    tree.add_right_child(node1, node3)
    tree.add_left_child(node2, node4)
    tree.add_right_child(node2, node5)
    tree.add_left_child(node3, 'G')
    tree.add_left_child(node4, 'H')
    tree.add_right_child(node4, 'I')
    tree.add_right_child(node5, 'J')
    tree.show_tree()

运行结果：

--------------------
A
├─R─C
| ├─R─F
| | ├─R─J
| | └─L─
| └─L─E
|   ├─R─I
|   └─L─H
└─L─B
  ├─R─D
  | ├─R─
  | └─L─G
  └─L─
--------------------

添加数据后的二叉树如下图：

四、二叉树的四种遍历方式

实现二叉树的层序遍历、先序遍历、中序遍历和后序遍历。

    def levelorder_traversal(self):
        """层序遍历"""
        if self.is_empty():
            print('空二叉树')
            return
        queue = list()
        queue.insert(0, self.__root)
        while len(queue):
            cur = queue.pop()
            print(cur.data, end=' ')
            if cur.left_child is not None:
                queue.insert(0, cur.left_child)
            if cur.right_child is not None:
                queue.insert(0, cur.right_child)
        print()

    def preorder_traversal(self, node):
        """先序遍历"""
        if node is None:
            return
        print(node.data, end=' ')
        self.preorder_traversal(node.left_child)
        self.preorder_traversal(node.right_child)

    def inorder_traversal(self, node):
        """中序遍历"""
        if node is None:
            return
        self.inorder_traversal(node.left_child)
        print(node.data, end=' ')
        self.inorder_traversal(node.right_child)

    def postorder_traversal(self, node):
        """后序遍历"""
        if node is None:
            return
        self.postorder_traversal(node.left_child)
        self.postorder_traversal(node.right_child)
        print(node.data, end=' ')

levelorder_traversal(): 二叉树的层序遍历。层序遍历是广度优先遍历，按从上到下、从左到右的顺序遍历二叉树。二叉树的层序遍历必须借助队列或栈，在 Python 中可以直接使用 list 来当队列使用。

preorder_traversal(): 二叉树的先序遍历。

inorder_traversal(): 二叉树的中序遍历。

postorder_traversal(): 二叉树的后序遍历。

关于二叉树的三种深度优先遍历方式，这里使用的是递归的方式，代码是很简单的，关键是理解遍历的顺序。

可以参考：https://blog.csdn.net/weixin_43790276/article/details/105473527

    print('层次遍历： ', end='')
    tree.levelorder_traversal()
    print('先序遍历： ', end='')
    tree.preorder_traversal(tree.root)
    print()
    print('中序遍历： ', end='')
    tree.inorder_traversal(tree.root)
    print()
    print('后序遍历： ', end='')
    tree.postorder_traversal(tree.root)
    print()

使用这四种方式遍历上面添加数据后的二叉树，运行结果如下：

层次遍历： A B C D E F G H I J 
先序遍历： A B D G C E H I F J 
中序遍历： B G D A H E I C F J 
后序遍历： G D B H I E J F C A 

五、二叉树的高度、叶节点和节点数

实现返回二叉树高度(深度)的方法，打印所有叶节点的方法，返回二叉树节点数量的方法，返回第 k 层节点数量的方法。

    def height(self, root):
        """二叉树的深度"""
        if root.data is None:
            return 0
        if root.left_child is None and root.right_child is None:
            return 1
        if root.left_child is not None and root.right_child is None:
            return 1 + self.height(root.left_child)
        if root.left_child is None and root.right_child is not None:
            return 1 + self.height(root.right_child)
        if root.left_child is not None and root.right_child is not None:
            return 1 + max(self.height(root.left_child), self.height(root.right_child))

    def leaves(self, root):
        """二叉树的叶节点"""
        if root.data is None:
            return None
        if root.left_child is None and root.right_child is None:
            print(root.data, end=' ')
        if root.left_child is not None:
            self.leaves(root.left_child)
        if root.right_child is not None:
            self.leaves(root.right_child)

    def node_count(self, root):
        """二叉树的节点个数"""
        return self.node_count(root.left_child) + self.node_count(root.right_child) + 1 if root else 0

    def kth_node_count(self, root, k):
        """二叉树第k层的节点个数"""
        if not root or k <= 0:
            return 0
        if k == 1:
            return 1
        return self.kth_node_count(root.left_child, k-1) + self.kth_node_count(root.right_child, k-1)

height(root): 返回二叉树的高度(深度)。如果二叉树是一棵空二叉树，则树的深度为 0。如果二叉树非空，只有根节点时，二叉树的深度为 1，层数每增加一层，二叉树的深度就加一，使用递归的方式一直找到最深一层。不管是左子树更深还是右子树更深，二叉树的深度都取两者中的最大值。

leaves(root): 获取二叉树的所有叶节点。二叉树的叶节点是没有子节点的节点，如果二叉树是一棵空二叉树，则没有叶节点。如果二叉树非空，只有根节点时，根节点本身就是叶节点。如果根节点有子树，使用递归的方式对子树进行判断，依次找到所有的叶节点。

node_count(root): 返回二叉树的节点个数。如果二叉树是一棵空二叉树，则节点个数为 0。如果二叉树非空，只有根节点时，节点个数为 1，再使用递归的方式对二叉树的左子树和右子树进行判断，节点个数为左子树加右子树加根节点的个数。

kth_node_count(root, k): 返回二叉树第 k 层的节点个数。如果二叉树是一棵空二叉树，则节点个数为 0。如果二叉树非空，k 小于等于 0 时，节点个数为 0，k 为 1 时，节点个数为 1，当 k 为其他值时，返回左子树和右子树的第 k-1 层的节点个数之和。这里为什么要减 1 呢？对于整棵二叉树来说，如 k=3 的节点在树的第三层，在左子树和右子树中，这些节点在左子树和右子树的第二层(3-1)，以此类推，递归就可以得出第 k 层的节点个数。

    print('二叉树的深度为：', tree.height(tree.root))
    print('二叉树的叶节点有：', end='')
    tree.leaves(tree.root)
    print()
    print('二叉树的节点数为：', tree.node_count(tree.root))
    print('二叉树的第三层节点数为：', tree.kth_node_count(tree.root, 3))

使用这几个方法查看上面添加节点后的二叉树，运行结果如下：

二叉树的深度为： 4
二叉树的叶节点有：G H I J 
二叉树的节点数为： 10
二叉树的第三层节点数为： 3

六、完整代码

# coding=utf-8
class Node(object):
    """节点类"""
    def __init__(self, data, left_child=None, right_child=None):
        self.data = data
        self.parent = None
        self.left_child = left_child
        self.right_child = right_child


class BinaryTree(object):
    """二叉树类"""
    def __init__(self):
        self.__root = None
        self.prefix_branch = '├'
        self.prefix_trunk = '|'
        self.prefix_leaf = '└'
        self.prefix_empty = ''
        self.prefix_left = '─L─'
        self.prefix_right = '─R─'

    def is_empty(self):
        return not self.__root

    @property
    def root(self):
        return self.__root

    @root.setter
    def root(self, data):
        self.__root = Node(data)

    def show_tree(self):
        if self.is_empty():
            print('空二叉树')
            return
        print('-' * 20)
        print(self.__root.data)
        self.__print_tree(self.__root)
        print('-' * 20)

    def add_left_child(self, parent, value):
        """添加左子节点"""
        if parent is None:
            print("父节点不存在")
            return
        if parent.left_child is not None:
            print("父节点已有左子节点")
            return
        node = value if isinstance(value, Node) else Node(value)
        parent.left_child = node
        node.parent = parent

    def add_right_child(self, parent, value):
        """添加右子节点"""
        if parent is None:
            print("父节点不存在")
            return
        if parent.right_child is not None:
            print("父节点已有右子节点")
            return
        node = value if isinstance(value, Node) else Node(value)
        parent.right_child = node
        node.parent = parent

    def levelorder_traversal(self):
        """层序遍历"""
        if self.is_empty():
            print('空二叉树')
            return
        queue = list()
        queue.insert(0, self.__root)
        while len(queue):
            cur = queue.pop()
            print(cur.data, end=' ')
            if cur.left_child is not None:
                queue.insert(0, cur.left_child)
            if cur.right_child is not None:
                queue.insert(0, cur.right_child)
        print()

    def preorder_traversal(self, node):
        """先序遍历"""
        if node is None:
            return
        print(node.data, end=' ')
        self.preorder_traversal(node.left_child)
        self.preorder_traversal(node.right_child)

    def inorder_traversal(self, node):
        """中序遍历"""
        if node is None:
            return
        self.inorder_traversal(node.left_child)
        print(node.data, end=' ')
        self.inorder_traversal(node.right_child)

    def postorder_traversal(self, node):
        """后序遍历"""
        if node is None:
            return
        self.postorder_traversal(node.left_child)
        self.postorder_traversal(node.right_child)
        print(node.data, end=' ')

    def height(self, root):
        """二叉树的深度"""
        if root.data is None:
            return 0
        if root.left_child is None and root.right_child is None:
            return 1
        if root.left_child is not None and root.right_child is None:
            return 1 + self.height(root.left_child)
        if root.left_child is None and root.right_child is not None:
            return 1 + self.height(root.right_child)
        if root.left_child is not None and root.right_child is not None:
            return 1 + max(self.height(root.left_child), self.height(root.right_child))

    def leaves(self, root):
        """二叉树的叶节点"""
        if root.data is None:
            return None
        if root.left_child is None and root.right_child is None:
            print(root.data, end=' ')
        if root.left_child is not None:
            self.leaves(root.left_child)
        if root.right_child is not None:
            self.leaves(root.right_child)

    def node_count(self, root):
        """二叉树的节点个数"""
        return self.node_count(root.left_child) + self.node_count(root.right_child) + 1 if root else 0

    def kth_node_count(self, root, k):
        """二叉树第k层的节点个数"""
        if not root or k <= 0:
            return 0
        if k == 1:
            return 1
        return self.kth_node_count(root.left_child, k-1) + self.kth_node_count(root.right_child, k-1)

    def __print_tree(self, node, prefix=None):
        if prefix is None:
            prefix, prefix_left_child = '', ''
        else:
            prefix = prefix.replace(self.prefix_branch, self.prefix_trunk)
            prefix = prefix.replace(self.prefix_leaf, self.prefix_empty)
            prefix_left_child = prefix.replace(self.prefix_leaf, self.prefix_empty)
        if self.has_child(node):
            if node.right_child is not None:
                print(prefix + self.prefix_branch + self.prefix_right + str(node.right_child.data))
                if self.has_child(node.right_child):
                    self.__print_tree(node.right_child, prefix + self.prefix_branch + ' ')
            else:
                print(prefix + self.prefix_branch + self.prefix_right)
            if node.left_child is not None:
                print(prefix + self.prefix_leaf + self.prefix_left + str(node.left_child.data))
                if self.has_child(node.left_child):
                    prefix_left_child += '  '
                    self.__print_tree(node.left_child, self.prefix_leaf + prefix_left_child)
            else:
                print(prefix + self.prefix_leaf + self.prefix_left)

    def has_child(self, node):
        return node.left_child is not None or node.right_child is not None

    def __str__(self):
        return str(self.__class__)


if __name__ == '__main__':
    tree = BinaryTree()
    tree.root = 'A'
    node1 = Node('B')
    tree.add_left_child(tree.root, node1)
    node2 = Node('C')
    tree.add_right_child(tree.root, node2)
    node3 = Node('D')
    node4 = Node('E')
    node5 = Node('F')
    tree.add_right_child(node1, node3)
    tree.add_left_child(node2, node4)
    tree.add_right_child(node2, node5)
    tree.add_left_child(node3, 'G')
    tree.add_left_child(node4, 'H')
    tree.add_right_child(node4, 'I')
    tree.add_right_child(node5, 'J')
    tree.show_tree()

    print('层次遍历： ', end='')
    tree.levelorder_traversal()
    print('先序遍历： ', end='')
    tree.preorder_traversal(tree.root)
    print()
    print('中序遍历： ', end='')
    tree.inorder_traversal(tree.root)
    print()
    print('后序遍历： ', end='')
    tree.postorder_traversal(tree.root)
    print()

    print('二叉树的深度为：', tree.height(tree.root))
    print('二叉树的叶节点有：', end='')
    tree.leaves(tree.root)
    print()
    print('二叉树的节点数为：', tree.node_count(tree.root))
    print('二叉树的第三层节点数为：', tree.kth_node_count(tree.root, 3))