rbtree.py

from enum import Enum, auto


class Color(Enum):
    RED = auto()
    BLACK = auto()


class Node:
    def __init__(self, key):
        self.key = key
        self.color = Color.RED
        self.parent = None
        self.left = None
        self.right = None


class RBTree:
    def __init__(self):
        self.nil = Node(None)
        # self.nil.parent = self.nil.left = self.nil.right = self.nil
        self.nil.color = Color.BLACK
        self.root = self.nil

    def __contains__(self, key):
        node = self.find_node(key)
        return node.key == key

    def __iter__(self):
        def visit(node):
            if node.left != self.nil:
                yield from visit(node.left)
            yield node.key
            if node.right != self.nil:
                yield from visit(node.right)

        if self.root != self.nil:
            yield from visit(self.root)

    def minimum(self, x):
        while x.left != self.nil:
            x = x.left
        return x

    def find_node(self, key):
        node = self.root
        while node != self.nil:
            if node.key == key:
                return node
            if key < node.key:
                node = node.left
            else:
                node = node.right
        return node

    def rotate_left(self, x):
        y = x.right  # yをxの右の子とする
        x.right = y.left  # yの左部分木をxの右部分木にする
        if y.left != self.nil:
            y.left.parent = x
        y.parent = x.parent  # xの親をyの親にする
        if x.parent == self.nil:  # xが根の時
            self.root = y
        elif x == x.parent.left:  # xが親の左部分木の時
            x.parent.left = y
        else:  # xが親の右部分木の時
            x.parent.right = y
        y.left = x  # xをyの右の子にする
        x.parent = y  # xの親をyにする

    def rotate_right(self, y):
        x = y.left
        y.left = x.right
        if x.right != self.nil:
            x.right.parent = y
        x.parent = y.parent
        if y.parent == self.nil:
            self.root = x
        elif y == y.parent.right:
            y.parent.right = x
        else:
            y.parent.left = x
        x.right = y
        y.parent = x

    def insert(self, key):
        z = Node(key)
        self.insert_node(z)

    def insert_node(self, z):
        y = self.nil
        # 根から挿入先を探索する
        x = self.root
        while x != self.nil:
            y = x
            if z.key < x.key:  # z,keyがx.keyより小さければxの左部分木を探索
                x = x.left
            else:  # z.keyがx.keyより大きければxの右部分木を探索
                x = x.right
        z.parent = y
        if y == self.nil:
            self.root = z
        elif z.key < y.key:
            y.left = z
        else:
            y.right = z
        z.left = self.nil
        z.right = self.nil
        z.color = Color.RED
        self.insert_fixup(z)

    def insert_fixup(self, z):
        while z.parent.color == Color.RED:
            if z.parent == z.parent.parent.left:
                y = z.parent.parent.right
                if y.color == Color.RED:  # case1: 叔父が赤
                    z.parent.color = Color.BLACK
                    y.color = Color.BLACK
                    z.parent.parent.color = Color.RED
                    z = z.parent.parent
                else:
                    if z == z.parent.right:  # case2: 叔父が黒で、zは右の子
                        # 左回転させて、case3に帰着させる
                        z = z.parent
                        self.rotate_left(z)
                    # case3: 叔父が黒で、zは左の子
                    z.parent.color = Color.BLACK
                    z.parent.parent.color = Color.RED
                    self.rotate_right(z.parent.parent)
            else:
                y = z.parent.parent.left
                if y.color == Color.RED:
                    z.parent.color = Color.BLACK
                    y.color = Color.BLACK
                    z.parent.parent.color = Color.RED
                    z = z.parent.parent
                else:
                    if z == z.parent.left:
                        z = z.parent
                        self.rotate_right(z)
                    z.parent.color = Color.BLACK
                    z.parent.parent.color = Color.RED
                    self.rotate_left(z.parent.parent)
        self.root.color = Color.BLACK

    def transplant(self, u, v):
        if u.parent == self.nil:
            self.root = v
        elif u == u.parent.left:
            u.parent.left = v
        else:
            u.parent.right = v
        v.parent = u.parent

    def delete(self, key):
        node = self.find_node(key)
        self.delete_node(node)

    def delete_node(self, z):
        y = z
        y_original_color = y.color
        if z.left == self.nil:  # case1: zの左の子がNIL
            x = z.right
            self.transplant(z, z.right)
        elif z.right == self.nil:  # case2 zの右の子がNIL
            x = z.left
            self.transplant(z, z.left)
        else:  # case3: zが２つの子を持つ場合
            y = self.minimum(z.right)  # zの右の部分木から次節点yを探す
            y_original_color = y.color
            x = y.right
            if y.parent != z:
                self.transplant(y, y.right)
                y.right = z.right
                y.right.parent = y
            self.transplant(z, y)
            y.left = z.left
            y.left.parent = y
            y.color = z.color
        if y_original_color == Color.BLACK:
            self.delete_fixup(x)

    def delete_fixup(self, x):
        while x != self.root and x.color == Color.BLACK:
            if x == x.parent.left:
                w = x.parent.right
                if w.color == Color.RED:
                    w.color = Color.BLACK
                    x.parent.color = Color.RED
                    self.rotate_left(x.parent)
                    w = x.parent.right
                if w.left.color == Color.BLACK and w.right.color == Color.BLACK:
                    w.color == Color.RED
                    x = x.parent
                else:
                    if w.right.color == Color.BLACK:
                        w.left.color = Color.BLACK
                        w.color = Color.RED
                        self.rotate_right(w)
                        w = x.parent.right
                    w.color = x.parent.color
                    x.parent.color = Color.BLACK
                    w.right.color = Color.BLACK
                    self.rotate_left(x.parent)
                    x = self.root
            else:
                w = x.parent.left
                if w.color == Color.RED:
                    w.color = Color.BLACK
                    x.parent.color = Color.RED
                    self.rotate_right(x.parent)
                    w = x.parent.left
                if w.right.color == Color.BLACK and w.left.color == Color.BLACK:
                    w.color == Color.RED
                    x = x.parent
                else:
                    if w.left.color == Color.BLACK:
                        w.right.color = Color.BLACK
                        w.color = Color.RED
                        self.rotate_left(w)
                        w = x.parent.left
                    w.color = x.parent.color
                    x.parent.color = Color.BLACK
                    w.left.color = Color.BLACK
                    self.rotate_right(x.parent)
                    x = self.root
        x.color = Color.BLACK