/
CHRedBlackTree.m
229 lines (200 loc) · 7.83 KB
/
CHRedBlackTree.m
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
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
//
// CHRedBlackTree.m
// CHDataStructures
//
// Copyright © 2008-2021, Quinn Taylor
// Copyright © 2002, Phillip Morelock
//
#import <CHDataStructures/CHRedBlackTree.h>
#import "CHAbstractBinarySearchTree_Internal.h"
#pragma mark C Functions for Optimized Operations
static CHBinaryTreeNode * rotateNodeWithLeftChild(CHBinaryTreeNode *node) {
CHBinaryTreeNode *leftChild = node->left;
node->left = leftChild->right;
leftChild->right = node;
node->color = kRED;
leftChild->color = kBLACK;
return leftChild;
}
static CHBinaryTreeNode * rotateNodeWithRightChild(CHBinaryTreeNode *node) {
CHBinaryTreeNode *rightChild = node->right;
node->right = rightChild->left;
rightChild->left = node;
node->color = kRED;
rightChild->color = kBLACK;
return rightChild;
}
static CHBinaryTreeNode * rotateObjectOnAncestor(id anObject, CHBinaryTreeNode *ancestor) {
if ([ancestor->object compare:anObject] == NSOrderedDescending) {
if ([ancestor->left->object compare:anObject] == NSOrderedDescending) {
ancestor->left = rotateNodeWithLeftChild(ancestor->left);
} else {
ancestor->left = rotateNodeWithRightChild(ancestor->left);
}
return ancestor->left;
} else {
if ([ancestor->right->object compare:anObject] == NSOrderedDescending) {
ancestor->right = rotateNodeWithLeftChild(ancestor->right);
} else {
ancestor->right = rotateNodeWithRightChild(ancestor->right);
}
return ancestor->right;
}
}
static CHBinaryTreeNode * singleRotation(CHBinaryTreeNode *node, BOOL goingRight) {
CHBinaryTreeNode *save = node->link[!goingRight];
node->link[!goingRight] = save->link[goingRight];
save->link[goingRight] = node;
node->color = kRED;
save->color = kBLACK;
return save;
}
static CHBinaryTreeNode * doubleRotation(CHBinaryTreeNode *node, BOOL goingRight) {
node->link[!goingRight] = singleRotation(node->link[!goingRight], !goingRight);
return singleRotation(node, goingRight);
}
#pragma mark -
@implementation CHRedBlackTree
// NOTE: The header and sentinel nodes are initialized to black (0) by default.
/*
Basically, as you walk down the tree to insert, if the present node has two red children, color it red and change the two children to black. If its parent is red, the tree must be rotated. (Just change the root's color back to black if you changed it). Returns without incrementing the count if the object already exists in the tree.
*/
- (void)addObject:(id)anObject {
CHRaiseInvalidArgumentExceptionIfNil(anObject);
++mutations;
CHBinaryTreeNode *current, *parent, *grandparent, *greatgrandparent;
greatgrandparent = grandparent = parent = current = header;
sentinel->object = anObject;
NSComparisonResult comparison;
while ((comparison = [current->object compare:anObject])) {
greatgrandparent = grandparent;
grandparent = parent;
parent = current;
current = current->link[comparison == NSOrderedAscending];
// Check for the bad case of red parent and red sibling of parent
if (current->left->color == kRED && current->right->color == kRED) {
// Simple red violation: resolve with color flip
current->color = kRED;
current->left->color = kBLACK;
current->right->color = kBLACK;
// Hard red violation: rotations necessary
if (parent->color == kRED) {
// BOOL lastWentRight = (grandparent->right == parent);
// greatgrandparent->link[greatgrandparent->right == grandparent]
// = (parent->link[lastWentRight])
// ? singleRotation(grandparent, !lastWentRight)
// : doubleRotation(grandparent, !lastWentRight);
grandparent->color = kRED;
if ([grandparent->object compare:anObject] != [parent->object compare:anObject]) {
parent = rotateObjectOnAncestor(anObject, grandparent);
}
current = rotateObjectOnAncestor(anObject, greatgrandparent);
current->color = kBLACK;
}
}
}
[anObject retain];
if (current != sentinel) {
// If an existing node matched, simply replace the existing value.
[current->object release];
current->object = anObject;
} else {
++count;
current = [self _createNodeWithObject:anObject];
parent->link[([parent->object compare:anObject] == NSOrderedAscending)] = current;
// one last reorientation check...
// Color flip
current->color = kRED;
current->left->color = kBLACK;
current->right->color = kBLACK;
// Fix red violation
if (parent->color == kRED) {
grandparent->color = kRED;
if ([grandparent->object compare:anObject] != [parent->object compare:anObject]) {
rotateObjectOnAncestor(anObject, grandparent);
}
current = rotateObjectOnAncestor(anObject, greatgrandparent);
current->color = kBLACK;
}
header->right->color = kBLACK; // Always reset root to black
}
}
/**
@param anObject The object to be removed from the tree.
@bug Performance decays exponentially (not linearly) when removing objects.
@todo Speed up red-black removal. The EternallyConfuzzled.com tutorial opts to push a red node down the tree using rotations and flips to avoid a nasty case of deleting a black node. This is almost certainly what causes the performance problems.
@see http://www.stanford.edu/~blp/avl/libavl.html/Deleting-from-an-RB-Tree.html
@see http://eternallyconfuzzled.com/tuts/datastructures/jsw_tut_rbtree.aspx
*/
- (void)removeObject:(id)anObject {
CHRaiseInvalidArgumentExceptionIfNil(anObject);
if (count == 0) {
return;
}
++mutations;
CHBinaryTreeNode *current, *parent, *grandparent;
parent = current = header;
CHBinaryTreeNode *found = NULL, *sibling;
sentinel->object = anObject;
NSComparisonResult comparison;
BOOL isGoingRight = YES, prevWentRight = YES;
while (current->link[isGoingRight] != sentinel) {
grandparent = parent;
parent = current;
current = current->link[isGoingRight];
comparison = [current->object compare:anObject];
prevWentRight = isGoingRight;
isGoingRight = (comparison != NSOrderedDescending);
if (comparison == NSOrderedSame) {
found = current; // Save a pointer; removal happens outside the loop
}
// There are only potential violations when removing a black node.
// If so, push the child red node down using rotations and color flips.
if (current->color != kRED && current->link[isGoingRight]->color != kRED) {
if (current->link[!isGoingRight]->color == kRED) {
parent->link[prevWentRight] = singleRotation(current, isGoingRight);
parent = parent->link[prevWentRight];
} else {
sibling = parent->link[prevWentRight];
if (sibling != sentinel) {
if (sibling->left->color == kBLACK && sibling->right->color == kBLACK) {
// If sibling's children are both black, do a color flip
parent->color = kBLACK;
sibling->color = kRED;
current->color = kRED;
} else {
CHBinaryTreeNode *tempNode = grandparent->link[(grandparent->right == parent)];
if (sibling->link[prevWentRight]->color == kRED) {
tempNode = doubleRotation(parent, prevWentRight);
} else if (sibling->link[!prevWentRight]->color == kRED) {
tempNode = singleRotation(parent, prevWentRight);
}
/* Ensure correct coloring */
current->color = tempNode->color = kRED;
tempNode->left->color = kBLACK;
tempNode->right->color = kBLACK;
}
} // if (sibling != sentinel)
}
}
}
// Transfer replacement value up to outgoing node, remove the "donor" node.
if (found != NULL) {
[found->object release];
found->object = current->object;
parent->link[(parent->right == current)]
= current->link[(current->left == sentinel)];
free(current);
--count;
}
header->right->color = kBLACK; // Make the root black for simplified logic
}
- (NSString *)debugDescriptionForNode:(CHBinaryTreeNode *)node {
return [NSString stringWithFormat:@"[%s]\t\"%@\"",
(node->color == kRED) ? " RED " : "BLACK", node->object];
}
- (NSString *)dotGraphStringForNode:(CHBinaryTreeNode *)node {
return [NSString stringWithFormat:@" \"%@\" [color=%@];\n",
node->object, (node->color == kRED) ? @"red" : @"black"];
}
@end