赤黒木

概要

平衡探索二分木の一種で、以下を満たすように木を保つ。

各ノードを赤と黒で（仮想的に）塗り分ける
根は黒
赤の子は黒（赤は2つ以上連続しない）
根からそれぞれの葉までの経路上にある黒の個数は等しい

こんなことを決めて何がいいのかと、

これが満たされている木では、根から最短経路（黒黒黒……）と最長経路（黒赤黒赤……）の差は2倍を超えないので、ほどよい平衡が保たれることが保証される
データを追加したり削除したりするとき、親子で繋がった直近3～4ノードの色のパターンを見て塗り替えることで状態を維持できる
- その中だけで状態を解消できないときは親に遡ることはあるが、再帰的に実装できる

ただ、パターン分けは結構種類があってややこしい。詳細は参考の1つめのサイトが丁寧に場合分けを説明されている。

参考
- Algorithms with Python / 赤黒木
- 赤黒木 - yaketake08's 実装メモ

実装

再帰関数は極力無くした実装（1箇所だけ、多重再帰にはならない箇所で使用）。どうしても配列アクセスは多くなるので、Pythonだと遅い。PyPyなら……速いとは言えないが、まぁまぁ。

関数	返値	概要
insert(x)	None	$x$ を追加（実装によっては重複する値を許さないこともできる。以下の実装では許す）
delete(x)	None	$x$ を削除（無ければ何もしない）
upper_bound(x)	$y, i$	$x < y$ を満たす最小の $y$ と、その小さい方からのindex $i$ (0-indexed)を取得
lower_bound(x)	$y, i$	$x \leq y$ を満たす最小の $y$ と、その小さい方からのindex $i$ (0-indexed)を取得
get_by_index(i)	$x$	小さい方から $i$ 番目の値を取得(0-indexed)
debug_print()	None	木を90°倒して左が根、上が小さい方となるように、“色,値,部分木以下の要素数”を描画

Verification: AtCoder Beginner Contest 140 E - Second Sum（PyPy3, 949ms）

"""
赤黒木
"""
 
 
class RedBlackTree:
    # For performance, nodes are implemented as lists.
    #   [left, right, data, color, count]
    #   color: 1:RED  0:BLACK
    #
    # It is useful to set the left/right index to 0/1 for boolean access.
    # The end nodes are BLACK and of which value is EOT (End Of Tree).
    #
    # Usage example:
    # rbt = RedBlackTree()
    # rbt.insert(x)
    # rbt.delete(x)
    # rbt.upper_bound(x)
    # rbt.lower_bound(x)
    # rbt.get_by_index(i)
    # rbt.debug_print()
     
 
    def __init__(self, EOT=-1):
        self.EOT= EOT
        self.root= self._leaf()
 
    def _leaf(self):
        return [None,None,self.EOT,0,0]
 
    def _rotate(self, node, r):
        child= node[r ^1]
        node[r ^1]= child[r]
        child[r]= node
        child[3]= node[3]
        node[3]= 1
        child_count= child[4]
        child[4]= node[4]
        node[4]-= child_count- node[r ^1][4]
        return child
 
    def insert(self, x):
        stack= []
        node= self.root
        while node[2] != self.EOT:
            # Variation: If the same value is not allowed, return when found.
 
            to_right= x >= node[2]
            stack.append((node, to_right))
            node= node[to_right]
 
        # Insert values to the end node
        node[0]= self._leaf()
        node[1]= self._leaf()
        node[2]= x
        node[3]= 1
        node[4]= 1
 
        # Increase count
        for parent, _in stack:
            parent[4]+= 1
 
        # Validate tree and rotate if needed.
        while stack:
            parent, r= stack.pop()
            if parent[3]== 1:
                node= parent
                continue
            parent[r]= node
            node, flag= self._balance_insert(parent, r)
 
            if stackand flag== True:
                parent, r= stack.pop()
                parent[r]= node
                break
 
        else:
            # If the root might have changed, update its color.
            self.root= node
            self.root[3]= 0
 
    def _balance_insert(self, node, r):
        flag= True
        if node[r][r ^1][3]== 1:
            node[r]= self._rotate(node[r], r)
        if node[r][r][3]== 1:
            if node[r ^1][3]== 1:
                node[3]= 1
                node[r ^1][3]= node[r][3]= 0
                flag= False
            else:
                node= self._rotate(node, r ^1)
        return node, flag
 
    def delete(self, x):
        node= self.root
        stack= []
        while node[2] != x:
            r= node[2] < x
            stack.append((node, r))
            node= node[r]
 
        # Not Found
        if node[2]== self.EOT:
            return
 
        # Node has 2 children: swap min node in right tree and delete min node
        if node[0][2] != self.EOTand node[1][2] != self.EOT:
            stack.append((node,1))
            min_node= self.get_min(node[1], stack)
            node[2]= min_node[2]
            node= min_node
 
        # Delete node is root
        if not stack:
            if node[0][2]== self.EOT:
                self.root= node[1]
            else:
                self.root= node[0]
            self.root[3]= 0
            return
 
        # Decrease count
        for parent, _in stack:
            parent[4]-= 1
 
        # Node has 0/1 child
        parent, r= stack[-1]
        if node[0][2]== self.EOT:
            parent[r]= node[1]
            node[1][3]= 0
            # Balance is only needed if both children are EOT and self color is black.
            if node[1][2] != self.EOTor node[3]== 1:
                return
        elif node[1][2]== self.EOT:
            parent[r]= node[0]
            node[0][3]= 0
            return
 
        # Validate tree and rotate if needed.
        while stack:
            parent, r= stack.pop()
            node, flag= self._balance_delete(parent, r)
 
            if stackand flag== True:
                parent, r= stack.pop()
                parent[r]= node
                break
        else:
            # If the root might have changed, update its color.
            self.root= node
            self.root[3]= 0
 
    def get_min(self, node, stack):
        while node[0][2] != self.EOT:
            stack.append((node,0))
            node= node[0]
        return node
 
    def _balance_delete(self, node, r):
        if node[r ^1][r][3]== 0 and node[r ^1][r ^1][3]== 0:
            if node[r ^1][3]== 0:
                node[r ^1][3]= 1
                if node[3]== 0:
                    return node,False
                node[3]= 0
            else:
                node= self._rotate(node, r)
                node[r], _= self._balance_delete(node[r], r)
        else:
            if node[r ^1][r][3]== 1:
                node[r ^1]= self._rotate(node[r ^1], r ^1)
            node= self._rotate(node, r)
            node[r][3]= 0
            node[r ^1][3]= 0
        return node,True
 
    def upper_bound(self, x):
        """
        :return Smallest y satisfying x < y and its leftmost index.
                If not exists, (EOT, length) will be returned.
        """
        node= self.root
        y= self.EOT
        i= node[4]
        j= 0
        while node[2] != self.EOT:
            if x < node[2]:
                y= node[2]
                i= j+ node[0][4]
                node= node[0]
            else:
                j+= node[0][4]+ 1
                node= node[1]
        return y, i
 
    def lower_bound(self, x):
        """
        :return Smallest y satisfying x <= y and its leftmost index.
                If not exists, (EOT, length) will be returned.
        """
        node= self.root
        y= self.EOT
        i= node[4]
        j= 0
        while node[2] != self.EOT:
            if x <= node[2]:
                y= node[2]
                i= j+ node[0][4]
                node= node[0]
            else:
                j+= node[0][4]+ 1
                node= node[1]
        return y, i
 
    def get_by_index(self, i):
        """
        :return (0-indexed) i-th smallest item.
                If i is greater than length, EOT will be returned.
        """
        node= self.root
        if node[4] <= i:
            return self.EOT
        j= i
        while node[2] != self.EOT:
            left_count= node[0][4]
            if left_count== j:
                return node[2]
            elif left_count > j:
                node= node[0]
            else:
                j-= left_count+ 1
                node= node[1]
 
    def debug_print(self):
        self._debug_print(self.root,0)
 
    def _debug_print(self, node, depth):
        if node[2] != self.EOT:
            self._debug_print(node[0], depth+ 1)
            print('      ' * depth,'BR'[node[3]], node[2], node[4])
            self._debug_print(node[1], depth+ 1)
 
 
def check():
    rbt= RedBlackTree()
    for xin [0,1,2,3,4,5,6,3,3,3,5,5,5,0,0,0]:
        print('------ INSERT', x,'------')
        rbt.insert(x)
        rbt.debug_print()
 
    for xin [0,1,2,3,4,5,6,7]:
        print('------ UPPER', x,'=', rbt.upper_bound(x))
        print('------ LOWER', x,'=', rbt.lower_bound(x))
 
    for xin [5,0,3,3,0,5]:
        print('------ DELETE', x,'------')
        rbt.delete(x)
        rbt.debug_print()
 
    for iin range(11):
        print('------ INDEX', i,'=', rbt.get_by_index(i))
 
 
if __name__== '__main__':
    check()