树结构的python实现（三：红黑树）

以下为红黑树的python实现：

from enum import Enum
from binary_search_tree import AVLTree, Node  # binary_search_tree的代码见上一篇文章：https://blog.youkuaiyun.com/moyao_miao/article/details/136827024

Color = Enum('Color', ('red', 'black'))  # 枚举节点的颜色（红或黑）
Direction = Enum('Direction', ('left', 'right'))  # 枚举节点的方向（左或右）


class RedBlackNode(Node):
    """红黑树节点类"""
    def __init__(self, value):
        super().__init__(value)
        self.parent_node = None  # 增加属性记录父节点
        self.color = Color.red  # 初始化节点的颜色为红色

    def add(self, left: bool, value):
        """
        重写add，在添加新节点后绑定父节点，并返回添加的子节点。
        :param left:是否加在左子节点上
        :param value:节点值
        :return:添加的子节点
        """
        if left:
            self.left_node = self.instantiate_myself(value)
            self.left_node.parent_node = self
            return self.left_node
        else:
            self.right_node = self.instantiate_myself(value)
            self.right_node.parent_node = self
            return self.right_node


class RedBlackTree(AVLTree):
    """
    红黑树类
    **改进建议：**
    1. **支持更广泛的输入**：考虑让__init__方法支持空列表输入，并在后续添加元素时建立树的结构，这样可以提供更灵活的类使用方式。
    2. **性能优化**：如果`data_list`较大，当前的初始化方法可能会导致性能问题，尤其是如果后者是通过连续插入实现的。考虑实现更高效的批量构建树的策略，比如先排序`data_list`然后通过分治法创建平衡的红黑树等。
    3. **代码复用**: 如果出现了大量重复的逻辑操作或模式，可以考虑将其抽取到工具类或基类中，以提高代码的复用性和可维护性。
    """
    def __init__(self, data_list):
        self.root = RedBlackNode(data_list[0])
        self.root.color = Color.black  # 创建根节点，并将其颜色设置为黑色
        self._list_to_binarytree(data_list)

    def rotate_left(self, root):
        """重写左旋转，设置颜色并修改父节点指向"""
        axis = root.right_node
        root.color = Color.red
        axis.color = Color.black
        root.right_node = axis.left_node
        if axis.left_node: axis.left_node.parent_node = root
        axis.left_node = root
        axis.parent_node = root.parent_node
        root.parent_node = axis
        return axis

    def rotate_right(self, root):
        """重写右旋转，设置颜色并修改父节点指向"""
        axis = root.left_node
        root.color = Color.red
        axis.color = Color.black
        root.left_node = axis.right_node
        if axis.right_node: axis.right_node.parent_node = root
        axis.right_node = root
        axis.parent_node = root.parent_node
        root.parent_node = axis
        return axis

    def get_direction(self, node, parent_node):
        """获取子节点在父节点中的方向"""
        if parent_node is None: return
        return Direction.left if parent_node.left_node == node else Direction.right

    def rebalance(self, node):
        """
        递归恢复平衡
        :param node: 添加的节点，或需要调整平衡的节点
        :return: 最终返回的是根节点
        """
        parent_node = node.parent_node
        if parent_node:
            grandpa_node = parent_node.parent_node
            if grandpa_node is None: return parent_node# 如果爷爷节点不存在，说明父节点是根节点，不需要重平衡，直接返回父节点
            if parent_node.color == Color.red:
                uncle_node = grandpa_node.right_node if grandpa_node.left_node == parent_node else grandpa_node.left_node  # 获取叔叔节点
                if uncle_node and uncle_node.color == Color.red:
                    # 父节点和叔叔节点都是红色，则把爷爷节点设为红色（若是根节点则不用变色），父节点和叔叔节点设为黑色
                    if grandpa_node is not self.root: grandpa_node.color = Color.red
                    parent_node.color = uncle_node.color = Color.black
                else:
                    # 确定插入节点相对于父节点的方向和父节点相对于爷爷节点的方向
                    node_direction = self.get_direction(node, parent_node)
                    parent_direction = self.get_direction(parent_node, grandpa_node)
                    great_grandpa_node = grandpa_node.parent_node
                    grandpa_direction = self.get_direction(grandpa_node, great_grandpa_node)
                    if parent_direction == Direction.left:
                        # 左左情况右旋，左右情况左右旋
                        grandpa_node = self.rotate_right(grandpa_node) if node_direction == Direction.left else self.rotate_left_right(grandpa_node)
                    else:
                        # 右右情况左旋，右左情况右左旋
                        grandpa_node = self.rotate_left(grandpa_node) if node_direction == Direction.right else self.rotate_right_left(grandpa_node)
                    if grandpa_direction:
                        # 将调整后的爷爷节点链接到曾祖节点上
                        if grandpa_direction == Direction.left: great_grandpa_node.left_node = grandpa_node
                        else: great_grandpa_node.right_node = grandpa_node
                grandpa_node = self.rebalance(grandpa_node)  # 递归向上调整平衡
            return grandpa_node
        else:return node

    def _add(self, node, value):
        """重写_add，在添加节点后恢复红黑树平衡"""
        added_node = None
        while not added_node:
            if value < node.value:
                if node.left_node: node = node.left_node
                else: added_node = node.add(True, value)
            else:
                if node.right_node: node = node.right_node
                else: added_node = node.add(False, value)
        node = self.rebalance(added_node)
        return node

    def plot(self):
        """重写plot，在黑色节点上加下划线后缀_作区分"""
        def binarytree_to_list(node, line=0):
            nonlocal tree_list
            if line > len(tree_list) - 1: tree_list.append([])
            if node is None:
                tree_list[line].append('N')
                return
            else:tree_list[line].append(str(node.value) + ('_' if node.color == Color.black else ''))
            binarytree_to_list(node.left_node, line + 1)
            binarytree_to_list(node.right_node, line + 1)
            line -= 1

        tree_list = []
        binarytree_to_list(self.root)
        l = len(tree_list) - 1
        for i, line_list in enumerate(tree_list[:-1]):
            print((' ' * (l - i)).join([' ' * (l - i)] + line_list))

    def add_node(self, data):
        """重写add_node，每添加一个节点画一次图"""
        self.root = self._add(self.root, data)
        print(f'添加{data}之后的红黑树：')
        self.plot()


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

    data_list = [3, 5, 8, 4, 6, 7, 9, 2, 1]
    print_format(RedBlackTree(data_list))

输出：

添加5之后的红黑树：
    3_
  N 5
添加8之后的红黑树：
    5_
  3 8
添加4之后的红黑树：
      5_
    3_  8_
  N 4 N N
添加6之后的红黑树：
      5_
    3_  8_
  N 4 6 N
添加7之后的红黑树：
      5_
    3_  7_
  N 4 6 8
添加9之后的红黑树：
        5_
      3_   7
    N  4  6_  8_
  N N N N N 9
添加2之后的红黑树：
        5_
      3_   7
    2  4  6_  8_
  N N N N N N N 9
添加1之后的红黑树：
        5_
      3   7
    2_  4_  6_  8_
  1 N N N N N N 9
中序遍历：
1	2	3	4	5	6	7	8	9	

红黑树图示：
        5_
      3   7
    2_  4_  6_  8_
  1 N N N N N N 9