Algoogle

Algorithm for Programming Contest

赤黒木(insert/delete)

基本情報


計算量  
insert
erase
空間
用途 set

N := 要素数

解説


平衡二分探索木として有名. std::setの実装にも使われている.
赤黒木とは以下の5つの性質を満たすような2分木.

  • 節点の色は赤か黒
  • 根は黒
  • 葉は黒
  • 赤の子は共に黒
  • 節点からその任意子孫の葉までの経路に含まれる黒節点の数は等しい

実装はinsert/deleteベースでされている(赤黒木(merge/split)).
根の親と葉としてnilを使う.
nilは1つ用意して流用する.
実装の関係上NULLとかにするとアクセスしようとして死ぬ.

ch[2]に左右の子を入れる.
left, rightとかでやるとさらに場合分けがエグい.

ノードに適当にデータをもたせればstd::mapっぽい感じになる.

競技プログラミング的な用途としてはこれを永続化して使うことがあるかないか(永続赤黒木).
そのままではstd::setの下位互換.

コード


(rbtree.cpp) download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
template<class T> class rbtree {
public:
    struct node {
        T key;
        int red;
        node *par, *ch[2]; // {left, right}
        node(T key):key(key),red(1) {
            par = ch[0] = ch[1] = NULL;
        }
    };

    unsigned sz;
    node *root, *nil;

    rbtree():sz(0) {
        nil = new node(-1);
        nil->red = 0;
        root = nil->par = nil->ch[0] = nil->ch[1] = nil;
    }

    unsigned size() { return sz; }

    void print(){ print(root,0,0); }
    void print(node *v, int dep, int lr) {
        if(v == nil) return;
        print(v->ch[1],dep+1,1);
        for(int i = 0; i < dep; i++) cerr << "  ";
        if(!lr) cerr << "--";
        else if(lr == 1) cerr << "「";
        else cerr << "L";
        if(v->red) cerr << "\x1b[31m";
        cerr << v->key << endl;
        cerr << "\x1b[0m";
        print(v->ch[0],dep+1,2);
    }

    void rotate(node *x, int d) {
        int e = d^1;
        node *y = x->ch[e];
        x->ch[e] = y->ch[d];
        if(y->ch[d] != nil) y->ch[d]->par = x;
        y->par = x->par;
        if(x->par == nil) root = y;
        else if(x == x->par->ch[d]) x->par->ch[d] = y;
        else x->par->ch[e] = y;
        y->ch[d] = x;
        x->par = y;
    }

    node *find(T key) {
        node *x = root;
        while(x != nil) {
            if(key != x->key) x = x->ch[key > x->key];
            else return x;
        }
        return x;
    }

    node *insert(T key) {
        node *x = root, *y = nil;
        while(x != nil) {
            y = x;
            if(key != x->key) x = x->ch[key > x->key];
            else return x;
        }
        sz++;
        node *z = new node(key);
        z->par = y;
        if(y == nil) root = z;
        else y->ch[z->key >= y->key] = z;
        z->ch[0] = z->ch[1] = nil;
        insert_update(z);
        nil->par = nil->ch[0] = nil->ch[1] = nil;
        return z;
    }

    void insert_update(node *z) {
        node *y;
        while(z->par->red) {
            int d = z->par == z->par->par->ch[0]? 0: 1, e = d^1;
            y = z->par->par->ch[e];
            if(y->red) {
                z->par->red = 0;
                y->red = 0;
                z->par->par->red = 1;
                z = z->par->par;
            }
            else if(z == z->par->ch[e]) {
                z = z->par;
                rotate(z,d);
            }
            else {
                z->par->red = 0;
                z->par->par->red = 1;
                rotate(z->par->par,e);
            }
        }
        root->red = 0;
    }

    void erase(T key) { erase(find(key)); }

    void erase(node *z) {
        if(z == nil) return;
        sz--;
        node *y = z, *x;
        int prevy = y->red;
        if(z->ch[0] == nil) {
            x = z->ch[1];
            transplant(z,z->ch[1]);
        }
        else if(z->ch[1] == nil) {
            x = z->ch[0];
            transplant(z,z->ch[0]);
        }
        else {
            y = minimum(z->ch[1]);
            prevy = y->red;
            x = y->ch[1];
            if(y->par == z) x->par = y;
            else {
                transplant(y,y->ch[1]);
                y->ch[1] = z->ch[1];
                y->ch[1]->par = y;
            }
            transplant(z,y);
            y->ch[0] = z->ch[0];
            y->ch[0]->par = y;
            y->red = z->red;
        }
        nil->par = nil->ch[0] = nil->ch[1] = nil;
        if(!prevy) erase_update(x);
        delete z;
        nil->par = nil->ch[0] = nil->ch[1] = nil;
    }

    void transplant(node *u, node *v) {
        if(u->par == nil) root = v;
        else if(u == u->par->ch[0]) u->par->ch[0] = v;
        else u->par->ch[1] = v;
        v->par = u->par;
    }

    node *minimum(node* x) {
        while(x->ch[0] != nil) x = x->ch[0];
        return x;
    }

    void erase_update(node *x) {
        node *w = nil;
        while(x != root and !x->red) {
            print();
            int d = x == x->par->ch[0]? 0: 1, e = d^1;
            w = x->par->ch[e];
            if(w->red) {
                w->red = 0;
                x->par->red = 1;
                rotate(x->par,d);
                w = x->par->ch[e];
            }
            else if(!w->ch[d]->red and !w->ch[e]->red) {
                w->red = 1;
                x = x->par;
            }
            else if(!w->ch[e]->red) {
                w->ch[d]->red = 0;
                w->red = 1;
                rotate(w,e);
                w = x->par->ch[e];
            }
            else {
                w->red = x->par->red;
                x->par->red = 0;
                w->ch[e]->red = 0;
                rotate(x->par,d);
                x = root;
            }
        }
        x->red = 0;
        print();
    }
};

Comments