平衡二叉树代码

以下是一个简单的平衡二叉树实现，包含了插入、删除和查找节点等操作：

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

class AVLTree:
    def __init__(self):
        self.root = None

    def insert(self, val):
        self.root = self._insert(self.root, val)

    def _insert(self, node, val):
        if not node:
            return TreeNode(val)
        elif val < node.val:
            node.left = self._insert(node.left, val)
        else:
            node.right = self._insert(node.right, val)

        node.height = 1 + max(self._height(node.left), self._height(node.right))
        balance = self._balance(node)

        # Left Left Case
        if balance > 1 and val < node.left.val:
            return self._rotate_right(node)

        # Right Right Case
        if balance < -1 and val > node.right.val:
            return self._rotate_left(node)

        # Left Right Case
        if balance > 1 and val > node.left.val:
            node.left = self._rotate_left(node.left)
            return self._rotate_right(node)

        # Right Left Case
        if balance < -1 and val < node.right.val:
            node.right = self._rotate_right(node.right)
            return self._rotate_left(node)

        return node

    def delete(self, val):
        self.root = self._delete(self.root, val)

    def _delete(self, node, val):
        if not node:
            return node
        elif val < node.val:
            node.left = self._delete(node.left, val)
        elif val > node.val:
            node.right = self._delete(node.right, val)
        else:
            if node.left is None:
                temp = node.right
                node = None
                return temp
            elif node.right is None:
                temp = node.left
                node = None
                return temp
            temp = self._min_value_node(node.right)
            node.val = temp.val
            node.right = self._delete(node.right, temp.val)

        if node is None:
            return node

        node.height = 1 + max(self._height(node.left), self._height(node.right))
        balance = self._balance(node)

        # Left Left Case
        if balance > 1 and self._balance(node.left) >= 0:
            return self._rotate_right(node)

        # Right Right Case
        if balance < -1 and self._balance(node.right) <= 0:
            return self._rotate_left(node)

        # Left Right Case
        if balance > 1 and self._balance(node.left) < 0:
            node.left = self._rotate_left(node.left)
            return self._rotate_right(node)

        # Right Left Case
        if balance < -1 and self._balance(node.right) > 0:
            node.right = self._rotate_right(node.right)
            return self._rotate_left(node)

        return node

    def search(self, val):
        return self._search(self.root, val)

    def _search(self, node, val):
        if not node:
            return False
        elif val == node.val:
            return True
        elif val < node.val:
            return self._search(node.left, val)
        else:
            return self._search(node.right, val)

    def _height(self, node):
        if not node:
            return 0
        return node.height

    def _balance(self, node):
        if not node:
            return 0
        return self._height(node.left) - self._height(node.right)

    def _rotate_right(self, node):
        new_root = node.left
        node.left = new_root.right
        new_root.right = node
        node.height = 1 + max(self._height(node.left), self._height(node.right))
        new_root.height = 1 + max(self._height(new_root.left), self._height(new_root.right))
        return new_root

    def _rotate_left(self, node):
        new_root = node.right
        node.right = new_root.left
        new_root.left = node
        node.height = 1 + max(self._height(node.left), self._height(node.right))
        new_root.height = 1 + max(self._height(new_root.left), self._height(new_root.right))
        return new_root

    def _min_value_node(self, node):
        current = node
        while current.left is not None:
            current = current.left
        return current

其中，`TreeNode` 表示二叉树的节点，包含了值、左子树、右子树以及节点的高度信息。`AVLTree` 是平衡二叉树的实现类，包含了插入、删除、查找节点等操作，以及 AVL 树维护平衡的相关方法 `_height`、`_balance`和 `_rotate_right`、`_rotate_left`。其中，`_insert` 和 `_delete` 方法都包含了对 AVL 树平衡的维护。