二叉树相关操作代码实现

class Node(object):
    """树节点"""

    def __init__(self, elem):
        self.elem = elem
        self.lchild = None
        self.rchild = None


class BinaryTree(object):
    """二叉树"""
    # 二叉树深度
    depth = 0

    def __init__(self):
        self.root = None

    def is_empty(self):
        """判断二叉树是否为空"""
        return self.root is None

    def add(self, elem):
        """添加二叉树节点"""
        node = Node(elem)
        # 如果根节点为空
        if self.is_empty():
            self.root = node
            return
        # 队列只有根节点
        queue = [self.root]
        while queue:
            # 取出队列的第一个元素
            cur_node = queue.pop(0)
            if cur_node.lchild is None:
                # 如果左子树为空
                cur_node.lchild = node
                return
            else:
                # 当前节点入队
                queue.append(cur_node.lchild)
            if cur_node.rchild is None:
                cur_node.rchild = node
                return
            else:
                queue.append(cur_node.rchild)

    def breadth_first_travel(self):
        """广度优先遍历"""
        if self.is_empty():
            print('二叉树为空')
            return
        queue = [self.root]
        while queue:
            cur_node = queue.pop(0)
            print(cur_node.elem, end=' ')
            if cur_node.lchild is not None:
                # 如果左子树不为空，当前节点入队
                queue.append(cur_node.lchild)
            if cur_node.rchild is not None:
                # 如果右子树不为空，当前节点入队
                queue.append(cur_node.rchild)

    """深度遍历有三种，前序，中序后序"""

    def pre_order(self, node):
        """先序遍历"""
        if self.is_empty():
            print('二叉树为空')
            return
        # 节点为空则返回
        if node is None:
            return
        print(node.elem, end=' ')
        self.pre_order(node.lchild)
        self.pre_order(node.rchild)

    def in_order(self, node):
        """中序遍历"""
        if self.is_empty():
            print('二叉树为空')
            return
        # 节点为空则返回
        if node is None:
            return
        self.in_order(node.lchild)
        print(node.elem, end=' ')
        self.in_order(node.rchild)

    def post_order(self, node):
        """后序遍历"""
        if self.is_empty():
            print('二叉树为空')
            return
        # 节点为空则返回
        if node is None:
            return
        self.post_order(node.lchild)
        self.post_order(node.rchild)
        print(node.elem, end=' ')

    def binary_tree_depth(self, node):
        """先序遍历求二叉树深度"""
        if self.is_empty():
            return 0
        # 那么树的深度为左子树的深度加1；同样，如果只有右子树没有左子树
        # 那么树的深度为右子树的深度加1。如果既有左子树也有右子树
        # 那该树的深度就是左子树和右子树的最大值加1
        # 节点为空则返回
        if node is None:
            return 0
        left = self.binary_tree_depth(node.lchild)
        right = self.binary_tree_depth(node.rchild)
        return max(left, right) + 1


if __name__ == "__main__":
    tree = BinaryTree()
    tree.add(0)
    tree.add(1)
    tree.add(2)
    tree.add(3)
    tree.add(4)
    tree.add(5)
    tree.add(6)
    tree.add(7)
    tree.add(8)
    tree.add(9)
    tree.breadth_first_travel()
    print(" ")
    print(tree.binary_tree_depth(tree.root))
    tree.pre_order(tree.root)
    print(" ")
    tree.in_order(tree.root)
    print(" ")
    tree.post_order(tree.root)
    print(" ")