Skip to content

Commit 2d7a151

Browse files
kirillinkirillin
committed
Add erase from rb tree
1 parent c9e7b86 commit 2d7a151

2 files changed

Lines changed: 298 additions & 15 deletions

File tree

‎include/mystd/some_trees/redblack_tree.hpp‎

Lines changed: 215 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ template <class Key, class Value, class Compare = std::less<Key>>
2727
class redblack_tree {
2828
using ValueType = std::pair<Key, Value>;
2929

30-
private:
30+
public:
3131
struct Node {
3232
ValueType data;
3333
Node* left = nullptr;
@@ -54,6 +54,27 @@ class redblack_tree {
5454
}
5555
}
5656

57+
Node* sibling(Node* n) {
58+
if (n == n->parent->left) {
59+
return n->parent->right;
60+
} else {
61+
return n->parent->left;
62+
}
63+
}
64+
65+
void transplant(Node* a, Node* b) {
66+
if (a->parent == nullptr) {
67+
root = b;
68+
} else if (a == a->parent->left) {
69+
a->parent->left = b;
70+
} else {
71+
a->parent->right = b;
72+
}
73+
if (b != nullptr) {
74+
b->parent = a->parent;
75+
}
76+
}
77+
5778
/// @brief
5879
/// g g
5980
/// |child |child
@@ -121,20 +142,21 @@ class redblack_tree {
121142
return pivot;
122143
}
123144

124-
/// @brief
145+
/// @brief Defenetly it is same as in libstdc++
146+
/// but without header node (to simplify logic)
125147
/// left triangle to line
126148
/// g -> g
127149
/// / \ -> / \
128-
/// p u -> p u
150+
/// p u -> p u
129151
/// \ -> /
130152
/// n -> n
131153
///
132154
/// right triangle to line
133155
/// g -> g
134156
/// / \ -> / \
135-
/// u p -> u p
157+
/// u p -> u p
136158
/// / -> \
137-
/// n -> n
159+
/// n -> n
138160
///
139161
///
140162
///
@@ -160,7 +182,7 @@ class redblack_tree {
160182
u->color = node_colors::BLACK;
161183
g->color = node_colors::RED;
162184
node = g;
163-
continue;
185+
break;
164186
}
165187

166188
// uncle is black or nullptr
@@ -199,6 +221,7 @@ class redblack_tree {
199221

200222
root->color = node_colors::BLACK;
201223
}
224+
root->color = node_colors::BLACK;
202225
}
203226

204227
void clear(Node* node) {
@@ -208,9 +231,134 @@ class redblack_tree {
208231
delete node;
209232
}
210233

211-
static Node* minimum(Node* n) {
212-
while (n && n->left) n = n->left;
213-
return n;
234+
static Node* search_min(Node* x) {
235+
if (!x) {
236+
return nullptr;
237+
}
238+
239+
while (x->left) {
240+
x = x->left;
241+
}
242+
return x;
243+
}
244+
245+
static Node* search_max(Node* x) {
246+
if (!x) {
247+
return nullptr;
248+
}
249+
250+
while (x->right) {
251+
x = x->right;
252+
}
253+
return x;
254+
}
255+
256+
static Node* successor(Node* x) {
257+
if (!x) {
258+
return nullptr;
259+
}
260+
261+
if (x->right) {
262+
return search_min(x->right);
263+
}
264+
265+
Node* y = x->parent;
266+
while (y && x == y->right) {
267+
x = y;
268+
y = y->parent;
269+
}
270+
return y;
271+
}
272+
273+
static Node* predecessor(Node* x) {
274+
if (!x) {
275+
return nullptr;
276+
}
277+
278+
if (x->left) {
279+
return search_max(x->left);
280+
}
281+
282+
Node* y = x->parent;
283+
while (y && x == y->left) {
284+
x = y;
285+
y = y->parent;
286+
}
287+
return y;
288+
}
289+
290+
/// @brief en wiki aproach
291+
void fixDelete(Node* x, Node* parent) {
292+
while (x != root && (x == nullptr || x->color == node_colors::BLACK)) {
293+
if (x == (parent ? parent->left : nullptr)) {
294+
Node* w = parent->right;
295+
296+
// case 3
297+
if (w->color == node_colors::RED) {
298+
w->color = node_colors::BLACK;
299+
parent->color = node_colors::RED;
300+
parent = rotate_left(parent);
301+
w = parent->right;
302+
}
303+
304+
if ((w->left == nullptr || w->left->color == node_colors::BLACK) &&
305+
(w->right == nullptr || w->right->color == node_colors::BLACK)) {
306+
// case 2
307+
w->color = node_colors::RED;
308+
x = parent;
309+
parent = parent->parent;
310+
} else {
311+
if (w->right == nullptr || w->right->color == node_colors::BLACK) {
312+
// case 5
313+
if (w->left) w->left->color = node_colors::BLACK;
314+
w->color = node_colors::RED;
315+
w = rotate_right(w);
316+
w = parent->right;
317+
}
318+
319+
// case 6
320+
w->color = parent->color;
321+
parent->color = node_colors::BLACK;
322+
if (w->right) w->right->color = node_colors::BLACK;
323+
parent = rotate_left(parent);
324+
x = root;
325+
}
326+
} else { // symetry cases
327+
Node* w = parent->left;
328+
329+
if (w->color == node_colors::RED) {
330+
// case 3
331+
w->color = node_colors::BLACK;
332+
parent->color = node_colors::RED;
333+
parent = rotate_right(parent);
334+
w = parent->left;
335+
}
336+
337+
if ((w->left == nullptr || w->left->color == node_colors::BLACK) &&
338+
(w->right == nullptr || w->right->color == node_colors::BLACK)) {
339+
// case 2
340+
w->color = node_colors::RED;
341+
x = parent;
342+
parent = parent->parent;
343+
} else {
344+
if (w->left == nullptr || w->left->color == node_colors::BLACK) {
345+
// case 5
346+
if (w->right) w->right->color = node_colors::BLACK;
347+
w->color = node_colors::RED;
348+
w = rotate_left(w);
349+
w = parent->left;
350+
}
351+
352+
// case 6
353+
w->color = parent->color;
354+
parent->color = node_colors::BLACK;
355+
if (w->left) w->left->color = node_colors::BLACK;
356+
parent = rotate_right(parent);
357+
x = root;
358+
}
359+
}
360+
}
361+
if (x) x->color = node_colors::BLACK;
214362
}
215363

216364
public:
@@ -227,7 +375,7 @@ class redblack_tree {
227375
iterator_basic& operator++() {
228376
if (!current) return *this;
229377
if (current->right) {
230-
current = minimum(current->right);
378+
current = search_min(current->right);
231379
} else {
232380
Node* p = current->parent;
233381
while (p && current == p->left) {
@@ -287,6 +435,63 @@ class redblack_tree {
287435
insert_and_rebalance(insert_left, new_node, parent);
288436
}
289437

438+
void erase(const Key& key) {
439+
Node* node = root;
440+
Node* parent = nullptr;
441+
442+
// search node to remove
443+
while (node && key != node->data.first) {
444+
parent = node;
445+
446+
if (cmp(key, node->data.first)) {
447+
node = node->left;
448+
} else {
449+
node = node->right;
450+
}
451+
}
452+
453+
// node is not found
454+
if (!node) {
455+
return;
456+
}
457+
458+
Node* y = node;
459+
Node* x = nullptr;
460+
node_colors y_original_color = y->color;
461+
462+
// both cildren
463+
if (node->left != nullptr && node->right != nullptr) {
464+
y = search_min(node->right);
465+
y_original_color = y->color;
466+
x = y->right;
467+
468+
if (y->parent == node) {
469+
if (x) x->parent = y;
470+
} else {
471+
transplant(y, y->right);
472+
y->right = node->right;
473+
y->right->parent = y;
474+
}
475+
476+
transplant(node, y);
477+
y->left = node->left;
478+
y->left->parent = y;
479+
y->color = node->color;
480+
481+
} else {
482+
// 2. <= 1 child
483+
x = (node->left != nullptr) ? node->left : node->right;
484+
transplant(node, x);
485+
}
486+
487+
delete node;
488+
489+
// remove black node -> black height changed, fix it
490+
if (y_original_color == node_colors::BLACK) {
491+
fixDelete(x, (x ? x->parent : nullptr));
492+
}
493+
}
494+
290495
iterator begin() { return iterator(minimum(root)); }
291496

292497
iterator end() { return iterator(nullptr); }

‎tests/unit_tests/some_trees/test_redblack_tree.cpp‎

Lines changed: 83 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -10,15 +10,93 @@ namespace my::testing {
1010
TEST(RedBlackTreeTest, InsertEraseContains) {
1111
my::redblack_tree<int, int> t;
1212

13-
for (size_t i = 1; i <= 10; ++i) {
13+
for (size_t i = 1; i <= 3; ++i) {
1414
t.insert({i, i * 10});
1515
}
1616

17-
// for (auto& [k, v] : t) {
18-
// std::cout << k << "\n";
19-
// }
20-
2117
t.print();
18+
t.erase(2);
19+
t.print();
20+
}
21+
22+
TEST(RedBlackTreeTest, InsertEraseSimple) {
23+
my::redblack_tree<int, int> t;
24+
t.insert({10, 10});
25+
t.insert({20, 20});
26+
t.insert({30, 30});
27+
28+
EXPECT_EQ(t.root->data.first, 20);
29+
EXPECT_EQ(t.root->left->data.first, 10);
30+
EXPECT_EQ(t.root->right->data.first, 30);
31+
32+
EXPECT_EQ(t.root->color, my::node_colors::BLACK);
33+
EXPECT_EQ(t.root->left->color, my::node_colors::RED);
34+
EXPECT_EQ(t.root->right->color, my::node_colors::RED);
35+
}
36+
37+
TEST(RedBlackTreeTest, BalanceRotateLeft) {
38+
my::redblack_tree<int, int> t;
39+
t.insert({10, 10});
40+
t.insert({20, 20});
41+
t.insert({30, 30});
42+
43+
EXPECT_EQ(t.root->data.first, 20);
44+
EXPECT_EQ(t.root->left->data.first, 10);
45+
EXPECT_EQ(t.root->right->data.first, 30);
46+
47+
EXPECT_EQ(t.root->color, my::node_colors::BLACK);
48+
EXPECT_EQ(t.root->left->color, my::node_colors::RED);
49+
EXPECT_EQ(t.root->right->color, my::node_colors::RED);
50+
}
51+
52+
TEST(RedBlackTreeTest, BalanceRotateRight) {
53+
my::redblack_tree<int, int> t;
54+
t.insert({10, 10});
55+
t.insert({5, 20});
56+
t.insert({1, 30});
57+
58+
EXPECT_EQ(t.root->data.first, 5);
59+
EXPECT_EQ(t.root->left->data.first, 1);
60+
EXPECT_EQ(t.root->right->data.first, 10);
61+
62+
EXPECT_EQ(t.root->color, my::node_colors::BLACK);
63+
EXPECT_EQ(t.root->left->color, my::node_colors::RED);
64+
EXPECT_EQ(t.root->right->color, my::node_colors::RED);
65+
}
66+
67+
TEST(RedBlackTreeTest, BalanceFromBook) {
68+
my::redblack_tree<int, int> t;
69+
t.insert({7, 10});
70+
t.insert({3, 10});
71+
t.insert({18, 10});
72+
t.insert({10, 10});
73+
t.insert({22, 10});
74+
t.insert({8, 10});
75+
t.insert({11, 10});
76+
t.insert({26, 10});
77+
78+
EXPECT_EQ(t.root->data.first, 7);
79+
EXPECT_EQ(t.root->left->data.first, 3);
80+
EXPECT_EQ(t.root->right->data.first, 18);
81+
82+
EXPECT_EQ(t.root->color, my::node_colors::BLACK);
83+
EXPECT_EQ(t.root->left->color, my::node_colors::BLACK);
84+
EXPECT_EQ(t.root->right->color, my::node_colors::RED);
85+
86+
EXPECT_EQ(t.root->right->left->data.first, 10);
87+
EXPECT_EQ(t.root->right->right->data.first, 22);
88+
89+
EXPECT_EQ(t.root->right->left->color, my::node_colors::BLACK);
90+
EXPECT_EQ(t.root->right->right->color, my::node_colors::BLACK);
91+
92+
EXPECT_EQ(t.root->right->left->left->data.first, 8);
93+
EXPECT_EQ(t.root->right->left->right->data.first, 11);
94+
95+
EXPECT_EQ(t.root->right->left->left->color, my::node_colors::RED);
96+
EXPECT_EQ(t.root->right->left->right->color, my::node_colors::RED);
97+
98+
EXPECT_EQ(t.root->right->right->right->data.first, 26);
99+
EXPECT_EQ(t.root->right->right->right->color, my::node_colors::RED);
22100
}
23101

24102
} // namespace my::testing

0 commit comments

Comments
 (0)