树结构的python实现（二：二叉查找树和AVL树）

以下为二叉查找树和AVL树的python实现：

from binary_tree import Node, BinaryTree  # binary_tree的代码见上一篇文章：https://blog.csdn.net/moyao_miao/article/details/136787981

class BinarySearchTree(BinaryTree):
    """二叉查找树类"""
    def __init__(self, data_list):
        self.root = Node(data_list[0])  # 创建根节点
        self._list_to_binarytree(data_list)  # 将剩余元素插入到二叉搜索树

    def _add(self, node, value):
        """
        递归添加节点：根据二叉搜索树的性质递归地将值插入到树中。
        :param node:当前节点。
        :param value:待插入的值。
        :return:插入值后的节点。
        """
        # 根据值与当前节点值的大小，递归地在左或右子树插入值
        if value < node.value:
            if node.left_node: node.left_node = self._add(node.left_node, value)
            else: node.add(True, value)
        else:
            if node.right_node: node.right_node = self._add(node.right_node, value)
            else: node.add(False, value)
        return node

    def add_node(self, data):
        """添加单个节点：外部接口，用于添加一个节点到二叉搜索树。"""
        self.root = self._add(self.root, data)

    def _list_to_binarytree(self, data_list):
        for data in data_list[1:]:
            self.add_node(data)


class AVLNode(Node):
    """"AVL树节点类,包含额外的高度信息用于维持树的平衡。"""
    def __init__(self, value):
        super().__init__(value)
        self.height = 0


class AVLTree(BinarySearchTree):
    """AVL树类"""
    def __init__(self, data_list):
        self.root = AVLNode(data_list[0])
        self._list_to_binarytree(data_list)

    def get_height(self, node):
        """获取节点高度，如果节点为空，则返回-1。"""
        return -1 if node is None else node.height

    def reset_height(self, node):
        """重置节点高度：根据节点的左右子节点的高度，重新计算并设置当前节点的高度。"""
        node.height = max(self.get_height(node.left_node), self.get_height(node.right_node)) + 1

    def rotate_left(self, root):
        """
        左旋转：对给定的树进行左旋转操作
        ### 改进建议
        1. **优化节点高度更新** - 当前的实现在每次旋转操作后都会直接调用 `reset_height` 更新节点的高度。虽然这保证了高度的正确性，但可能会引入不必要的效率问题。考虑在旋转操作中直接更新高度，而不是通过重复调用 `reset_height` 方法。
        """
        axis = root.right_node  # 以树根的右子节点为轴
        root.right_node = axis.left_node  # 先把轴的左子节点设为根的右子节点
        axis.left_node = root  # 再把根设为轴的左子节点
        self.reset_height(root)
        self.reset_height(axis)
        return axis

    def rotate_right(self, root):
        """右旋转"""
        axis = root.left_node
        root.left_node = axis.right_node
        axis.right_node = root
        self.reset_height(root)
        self.reset_height(axis)
        return axis

    def rotate_left_right(self, root):
        """先左旋后右旋"""
        root.left_node = self.rotate_left(root.left_node)  # 以树根的左子节点的右子节点为轴
        return self.rotate_right(root)

    def rotate_right_left(self, root):
        """先右旋后左旋"""
        root.right_node = self.rotate_right(root.right_node)
        return self.rotate_left(root)

    def _add(self, node, value):
        """递归添加节点，并且在必要时通过旋转操作调整树的结构保持AVL树的平衡性质。"""
        if value < node.value:
            if node.left_node: node.left_node = self._add(node.left_node, value)
            else: node.add(True, value)
            # 检查并调整左侧失衡
            if self.get_height(node.left_node) - self.get_height(node.right_node) > 1:
                node = self.rotate_right(node) if value < node.left_node.value else self.rotate_left_right(node)
        else:
            if node.right_node: node.right_node = self._add(node.right_node, value)
            else: node.add(False, value)
            # 检查并调整右侧失衡
            if self.get_height(node.right_node) - self.get_height(node.left_node) > 1:
                node = self.rotate_left(node) if value > node.right_node.value else self.rotate_right_left(node)
        self.reset_height(node)  # 更新当前节点的高度
        return node


if __name__ == "__main__":
    def print_format(obj):
        print("中序遍历：");obj.inorder_traversal()
        print('二叉树图示：');obj.plot()

    data_list = list(reversed(range(10)))
    print_format(BinarySearchTree(data_list))
    print_format(AVLTree(data_list))

输出：

中序遍历：
0	1	2	3	4	5	6	7	8	9	

二叉树图示：
                    9
                  8         N
                7        N
              6       N
            5      N
          4     N
        3    N
      2   N
    1  N
  0 N
中序遍历：
0	1	2	3	4	5	6	7	8	9	

二叉树图示：
        6
      2   8
    1  4  7  9
  0 N 3 5 N N N N