Python中的AVL树实现

2023-04-11 00:00:00 python AVL

AVL树是一种自平衡二叉搜索树,在插入或删除节点时会进行旋转操作来保持树的平衡性。下面是Python中的AVL树实现:

class AVLNode:
    def __init__(self, key: str):
        self.key = key
        self.left = None
        self.right = None
        self.height = 1

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

    def insert(self, key: str):
        self.root = self._insert(self.root, key)

    def delete(self, key: str):
        self.root = self._delete(self.root, key)

    def _get_height(self, node: AVLNode) -> int:
        if not node:
            return 0
        return node.height

    def _get_balance_factor(self, node: AVLNode) -> int:
        if not node:
            return 0
        return self._get_height(node.left) - self._get_height(node.right)

    def _update_height(self, node: AVLNode):
        node.height = 1 + max(self._get_height(node.left), self._get_height(node.right))

    def _rotate_left(self, node: AVLNode):
        right_child = node.right
        right_left_child = right_child.left

        right_child.left = node
        node.right = right_left_child

        self._update_height(node)
        self._update_height(right_child)

        return right_child

    def _rotate_right(self, node: AVLNode):
        left_child = node.left
        left_right_child = left_child.right

        left_child.right = node
        node.left = left_right_child

        self._update_height(node)
        self._update_height(left_child)

        return left_child

    def _insert(self, node: AVLNode, key: str) -> AVLNode:
        if not node:
            return AVLNode(key)
        elif key < node.key:
            node.left = self._insert(node.left, key)
        else:
            node.right = self._insert(node.right, key)

        self._update_height(node)

        balance_factor = self._get_balance_factor(node)

        if balance_factor > 1 and key < node.left.key:
            return self._rotate_right(node)

        if balance_factor > 1 and key > node.left.key:
            node.left = self._rotate_left(node.left)
            return self._rotate_right(node)

        if balance_factor < -1 and key > node.right.key:
            return self._rotate_left(node)

        if balance_factor < -1 and key < node.right.key:
            node.right = self._rotate_right(node.right)
            return self._rotate_left(node)

        return node

    def _get_min_node(self, node: AVLNode) -> AVLNode:
        while node.left:
            node = node.left
        return node

    def _delete(self, node: AVLNode, key: str) -> AVLNode:
        if not node:
            return None

        elif key < node.key:
            node.left = self._delete(node.left, key)

        elif key > node.key:
            node.right = self._delete(node.right, key)

        else:
            if not node.left:
                temp = node.right
                node = None
                return temp

            elif not node.right:
                temp = node.left
                node = None
                return temp

            temp = self._get_min_node(node.right)
            node.key = temp.key
            node.right = self._delete(node.right, temp.key)

        if not node:
            return node

        self._update_height(node)

        balance_factor = self._get_balance_factor(node)

        if balance_factor > 1 and self._get_balance_factor(node.left) >= 0:
            return self._rotate_right(node)

        if balance_factor > 1 and self._get_balance_factor(node.left) < 0:
            node.left = self._rotate_left(node.left)
            return self._rotate_right(node)

        if balance_factor < -1 and self._get_balance_factor(node.right) <= 0:
            return self._rotate_left(node)

        if balance_factor < -1 and self._get_balance_factor(node.right) > 0:
            node.right = self._rotate_right(node.right)
            return self._rotate_left(node)

        return node

在以上代码中,AVLNode表示AVL树中的节点,含有节点的key值、左右子节点和节点的高度。AVLTree表示AVL树,含有树的根节点以及插入、删除、旋转等操作。其中,_get_height函数获取节点的高度,_get_balance_factor函数获取节点的平衡因子,_update_height函数更新节点的高度,_rotate_left函数向左旋转节点,_rotate_right函数向右旋转节点,_insert函数插入节点,_get_min_node函数获取最小节点,_delete函数删除节点。

以上是Python中的AVL树实现,可以用AVLTree类来创建AVL树对象并进行操作,例如:

tree = AVLTree()
tree.insert('pidancode.com')
tree.insert('Python')
tree.insert('AVL')
tree.delete('Python')

上面的代码创建了一个AVL树,插入了三个节点,然后删除了其中一个节点。

相关文章