首页 编写策略 数据平台 因子研究 我的交易 发现策略 量化学院 知识库 免费注册
AI量化知识树

AVL树和红黑树的Python代码实现

由bqw9z8tc创建，最终由small_q 被浏览 47 用户

AVL树

AVL树是一种自平衡二叉搜索树。在这种树中，任何节点的两个子树的高度差被严格控制在1以内。这确保了树的平衡，从而保证了搜索、插入和删除操作的高效性。AVL树是由Georgy Adelson-Velsky和Evgenii Landis在1962年发明的，因此得名（Adelson-Velsky和Landis树）。

平衡因子：每个节点的平衡因子是其左子树的高度减去其右子树的高度。平衡因子必须保持在-1、0或1之间。

旋转操作：为了维持平衡，在进行插入或删除操作后，可能会执行单旋转或双旋转。单旋转包括右旋（针对左重失衡）和左旋（针对右重失衡）。双旋转是一种更复杂的调整，包括左-右旋转和右-左旋转。

时间复杂度：在AVL树中，查找、插入和删除操作的时间复杂度均为O(log n)，其中n是树中的节点数。

AVL树节点定义

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

这里定义了一个AVL树的节点。每个节点有一个键值key，一个高度属性height，以及指向左右子节点的指针left和right。


AVL树的高度和平衡因子

AVL树的关键是保持每个节点的平衡。平衡因子被定义为节点的左子树高度减去其右子树高度。这个因子应该是-1、0或1。如果在任何时候，一个节点的平衡因子不在这个范围内，我们需要通过旋转来重新平衡树。

def get_height(node):
    if not node:
        return 0
    return node.height

def get_balance(node):
    if not node:
        return 0
    return get_height(node.left) - get_height(node.right)

这两个函数帮助我们计算任何节点的高度和平衡因子。

AVL树旋转操作

为了维持平衡，我们可能需要进行树的旋转操作。主要有四种类型的旋转：右旋转、左旋转、左-右旋转和右-左旋转。

  1. 右旋转：

当一个节点的左子树比右子树高时，我们进行右旋转。

def rotate_right(z):
    y = z.left
    T3 = y.right
    y.right = z
    z.left = T3
    z.height = 1 + max(get_height(z.left), get_height(z.right))
    y.height = 1 + max(get_height(y.left), get_height(y.right))
    return y
  1. 左旋转：

当一个节点的右子树比左子树高时，我们进行左旋转。

def rotate_left(z):
    y = z.right
    T2 = y.left
    y.left = z
    z.right = T2
    z.height = 1 + max(get_height(z.left), get_height(z.right))
    y.height = 1 + max(get_height(y.left), get_height(y.right))
    return y
  1. 左-右旋转和右-左旋转：

这些是双旋转操作，先进行一次旋转使其成为第一种或第二种情况，然后再进行相应的旋转。

AVL树插入操作

插入操作类似于普通的二叉搜索树插入，但是在插入后，我们需要更新祖先节点的高度，并检查每个节点的平衡因子，以确保树保持平衡。

def insert(node, key):
    if not node:
        return AVLNode(key)
    elif key < node.key:
        node.left = insert(node.left, key)
    else:
        node.right = insert(node.right, key)

    node.height = 1 + max(get_height(node.left), get_height(node.right))
    balance = get_balance(node)

    if balance > 1 and key < node.left.key:
        return rotate_right(node)
    if balance < -1 and key > node.right.key:
        return rotate_left(node)
    if balance > 1 and key > node.left.key:
        node.left = rotate_left(node.left)
        return rotate_right(node)
    if balance < -1 and key < node.right.key:
        node.right = rotate_right(node.right)
        return rotate_left(node)

    return node

AVL树删除操作

删除操作也与普通的二叉搜索树类似，但在删除节点后，我们需要更新祖先节点的高度，并检查每个节点的平衡因子以重新平衡树。

def min_value_node(node):
    current = node
    while current.left is not None:
        current = current.left
    return current

def delete(node, key):
    if not node:
        return node

    if key < node.key:
        node.left = delete(node.left, key)
    elif key > node.key:
        node.right = delete(node.right, key)
    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 = min_value_node(node.right)
        node.key = temp.key
        node.right = delete(node.right, temp.key)

    if node is None:
        return node

    node.height = 1 + max(get_height(node.left), get_height(node.right))
    balance = get_balance(node)

    if balance > 1 and get_balance(node.left) >= 0:
        return rotate_right(node)
    if balance < -1 and get_balance(node.right) <= 0:
        return rotate_left(node)
    if balance > 1 and get_balance(node.left) < 0:
        node.left = rotate_left(node.left)
        return rotate_right(node)
    if balance < -1 and get_balance(node.right) > 0:
        node.right = rotate_right(node.right)
        return rotate_left(node)

    return node

这个删除函数处理了三种情况：要删除的节点有两个子节点、一个子节点或没有子节点。在删除节点后，它会通过旋转操作确保树保持平衡。

使用AVL树

现在我们可以创建一个AVL树的实例，并进行插入和删除操作：

avl_tree = None

keys = [20, 4, 15, 70, 50, 80, 100]
for key in keys:
    avl_tree = insert(avl_tree, key)

avl_tree = delete(avl_tree, 70)

在这个例子中，我们插入了一些数字，然后删除了一个。

红黑树

红黑树是另一种自平衡二叉搜索树，它通过确保树从根到叶子的最长可能路径不超过最短可能路径的两倍来维持大致的平衡。红黑树的节点有两种颜色：红色或黑色。这种树遵循以下重要属性：

颜色属性：每个节点要么是红色，要么是黑色。

根属性：树的根总是黑色。

叶子属性：所有叶子（NIL节点）都是黑色。

红色节点属性：如果一个节点是红色的，那么它的两个子节点都是黑色的。

路径属性：从任一节点到其每个叶子的所有路径都包含相同数量的黑色节点。

红黑树通过旋转和重新着色来维持这些属性。虽然红黑树的平衡性不如AVL树，但它在插入和删除操作中需要更少的旋转，这使得它在某些类型的应用中更高效。

红黑树节点定义：

class RedBlackNode:
    def __init__(self, key, color="red"):
        self.key = key
        self.color = color
        self.left = None
        self.right = None
        self.parent = None

这里定义了一个红黑树的节点。除了常规的key外，节点还包含一个color属性以及指向父节点和子节点的链接。

红黑树旋转操作：

为了维持红黑树的特性，在插入和删除节点时可能需要进行树的旋转操作。主要有两种类型的旋转：左旋和右旋。

  1. 左旋：

当一个节点的右子节点是红色而左子节点是黑色时进行。

def rotate_left(self, x):
    y = x.right
    x.right = y.left
    if y.left != self.TNULL:
        y.left.parent = x

    y.parent = x.parent
    if x.parent == None:
        self.root = y
    elif x == x.parent.left:
        x.parent.left = y
    else:
        x.parent.right = y
    y.left = x
    x.parent = y
  1. 右旋：

当一个节点的左子节点是红色而左子节点的左子节点也是红色时进行。

def rotate_right(self, x):
    y = x.left
    x.left = y.right
    if y.right != self.TNULL:
        y.right.parent = x

    y.parent = x.parent
    if x.parent == None:
        self.root = y
    elif x == x.parent.right:
        x.parent.right = y
    else:
        x.parent.left = y
    y.right = x
    x.parent = y

红黑树插入操作：

插入操作首先类似于普通的二叉搜索树插入。但在插入一个节点后，可能需要进行多次颜色更改和旋转来修复可能违反的红黑树性质。

def insert(self, key):
    node = RedBlackNode(key)
    node.parent = None
    node.key = key
    node.left = self.TNULL
    node.right = self.TNULL
    node.color = 1  # 新节点必须是红色

    y = None
    x = self.root

    while x != self.TNULL:
        y = x
        if node.key < x.key:
            x = x.left
        else:
            x = x.right

    node.parent = y
    if y == None:
        self.root = node
    elif node.key < y.key:
        y.left = node
    else:
        y.right = node

    if node.parent == None:
        node.color = 0
        return

    if node.parent.parent == None:
        return

    self.fix_insert(node)

在插入节点后，fix_insert函数被调用以重新平衡树，确保树仍然是一个有效的红黑树。

红黑树删除操作：

删除操作比插入操作更复杂。删除节点后，可能需要进行多次颜色更改和旋转来修复可能违反的红黑树性质。

def delete_node(self, node, key):
    # 省略具体删除逻辑
    self.fix_delete(x)

在删除节点后，fix_delete函数被调用以重新平衡树。

使用红黑树：

现在我们可以创建一个红黑树的实例，并进行插入和删除操作：

rbt = RedBlackTree()
numbers = [20, 15, 25, 10, 30, 5, 35]
for number in numbers:
    rbt.insert(number)

rbt.delete_node(rbt.root, 15)

在这个例子中，我们插入了一些数字，然后删除了一个。

红黑树是一种复杂的数据结构，它通过一系列精心设计的规则和操作来确保树的平衡。尽管它的实现比普通的二叉搜索树复杂得多，但它在插入、删除和查找操作上提供了良好的最坏情况性能。这种平衡是通过确保任何路径上的黑色节点数目大致相同来实现的。红黑树广泛应用于计算机科学的许多领域，尤其是在那些需要快速查找操作的场景中，如数据库和文件系统。

实现红黑树要求对它的性质和操作有深入的理解。上述代码提供了一个基本的框架，但它并不完整。在实际应用中，您可能还需要处理更多的边界情况和优化。每个操作背后都有复杂的逻辑和数学原理，这是理解和实现红黑树的关键部分。

标签

python开发Python
{link}