update options for pip installation in README.md
[liba.git] / src / rbt.c
blob39f6c19ceaaa4ec3a98c2a84625de741ef5f256f
1 #include "a/rbt.h"
3 /*
4 red-black trees properties: https://en.wikipedia.org/wiki/Rbtree
6 1) A node is either red or black
7 2) The root is black
8 3) All leaves (null) are black
9 4) Both children of every red node are black
10 5) Every simple path from root to leaves contains the same number of black nodes.
12 4 and 5 give the O(log n) guarantee, since 4 implies you cannot have two
13 consecutive red nodes in a path and every red node is therefore followed by
14 a black. So if B is the number of black nodes on every simple path (as per
15 5), then the longest possible path due to 4 is 2B.
17 We shall indicate color with case, where black nodes are uppercase and red
18 nodes will be lowercase. Unknown color nodes shall be drawn as red within
19 parentheses and have some accompanying text comment.
22 /* Replaces the child of the specified red–black tree node. */
23 static A_INLINE void a_rbt_new_child(a_rbt *root, a_rbt_node *parent, a_rbt_node *node, a_rbt_node *newnode)
25 if (parent)
27 if (parent->left == node)
29 parent->left = newnode;
31 else
33 parent->right = newnode;
36 else
38 root->node = newnode;
42 /* Sets the parent and color of the specified red–black tree node. */
43 static A_INLINE void a_rbt_set_parent_color(a_rbt_node *node, a_rbt_node *parent, unsigned int color)
45 #if defined(A_SIZE_POINTER) && (A_SIZE_POINTER + 0 > 1)
46 node->parent_ = (a_uptr)parent | color;
47 #else /* !A_SIZE_POINTER */
48 node->parent = parent;
49 node->color = color;
50 #endif /* A_SIZE_POINTER */
53 /* Sets the parent of the specified red–black tree node. */
54 static A_INLINE void a_rbt_set_parent(a_rbt_node *node, a_rbt_node *parent)
56 #if defined(A_SIZE_POINTER) && (A_SIZE_POINTER + 0 > 1)
57 node->parent_ = (a_uptr)parent | (node->parent_ & 1);
58 #else /* !A_SIZE_POINTER */
59 node->parent = parent;
60 #endif /* A_SIZE_POINTER */
63 /* Returns the color of the specified red–black tree node. */
64 static A_INLINE unsigned int a_rbt_color(a_rbt_node const *node)
66 #if defined(A_SIZE_POINTER) && (A_SIZE_POINTER + 0 > 1)
67 return (unsigned int)(node->parent_ & 1);
68 #else /* !A_SIZE_POINTER */
69 return node->color;
70 #endif /* A_SIZE_POINTER */
73 /* Sets the black color of the specified red–black tree node. */
74 static A_INLINE void a_rbt_set_black(a_rbt_node *node)
76 #if defined(A_SIZE_POINTER) && (A_SIZE_POINTER + 0 > 1)
77 node->parent_ |= A_RBT_B;
78 #else /* !A_SIZE_POINTER */
79 node->color = A_RBT_B;
80 #endif /* A_SIZE_POINTER */
84 Helper function for rotations:
85 - node's parent and color get assigned to new
86 - node gets assigned new as a parent and 'color' as a color.
88 static A_INLINE void a_rbt_set_parents(a_rbt *root, a_rbt_node *node, a_rbt_node *newnode, unsigned int color)
90 a_rbt_node *const parent = a_rbt_parent(node);
91 #if defined(A_SIZE_POINTER) && (A_SIZE_POINTER + 0 > 1)
92 newnode->parent_ = node->parent_;
93 #else /* !A_SIZE_POINTER */
94 newnode->parent = node->parent;
95 newnode->color = node->color;
96 #endif /* A_SIZE_POINTER */
97 a_rbt_set_parent_color(node, newnode, color);
98 a_rbt_new_child(root, parent, node, newnode);
101 void a_rbt_insert_adjust(a_rbt *root, a_rbt_node *node)
103 for (a_rbt_node *parent = a_rbt_parent(node), *gparent, *tmp;;)
105 /* Loop invariant: node is red. */
106 if (a_unlikely(!parent))
109 The inserted node is root. Either this is the first node,
110 or we recursed at Case 1 below and are no longer violating 4).
112 a_rbt_set_parent_color(node, A_NULL, A_RBT_B);
113 break;
116 If there is a black parent, we are done. Otherwise, take some corrective action as, per 4),
117 we don't want a red root or two consecutive red nodes.
119 if (a_rbt_color(parent) == A_RBT_B) { break; }
120 gparent = a_rbt_parent(parent);
121 tmp = gparent->right;
122 if (parent != tmp) /* parent == gparent->left */
124 if (tmp && a_rbt_color(tmp) == A_RBT_R)
127 Case 1 - node's uncle is red (color flips).
130 / \ / \
131 p u --> P U
135 However, since g's parent might be red, and 4) does not allow this, we need to recurse at g.
137 a_rbt_set_parent_color(tmp, gparent, A_RBT_B);
138 a_rbt_set_parent_color(parent, gparent, A_RBT_B);
139 node = gparent;
140 parent = a_rbt_parent(node);
141 a_rbt_set_parent_color(node, parent, A_RBT_R);
142 continue;
144 tmp = parent->right;
145 if (node == tmp)
148 Case 2 - node's uncle is black and node is the parent's right child (left rotate at parent).
151 / \ / \
152 p U --> n U
156 This still leaves us in violation of 4), the continuation into Case 3 will fix that.
158 tmp = node->left;
159 parent->right = tmp;
160 node->left = parent;
161 if (tmp) { a_rbt_set_parent_color(tmp, parent, A_RBT_B); }
162 a_rbt_set_parent_color(parent, node, A_RBT_R);
163 parent = node;
164 tmp = node->right;
167 Case 3 - node's uncle is black and node is
168 the parent's left child (right rotate at gparent).
171 / \ / \
172 p U --> n g
176 gparent->left = tmp; /* == parent->right */
177 parent->right = gparent;
178 if (tmp) { a_rbt_set_parent_color(tmp, gparent, A_RBT_B); }
179 a_rbt_set_parents(root, gparent, parent, A_RBT_R);
180 break;
182 else
184 tmp = gparent->left;
185 if (tmp && a_rbt_color(tmp) == A_RBT_R)
187 /* Case 1 - color flips */
188 a_rbt_set_parent_color(tmp, gparent, A_RBT_B);
189 a_rbt_set_parent_color(parent, gparent, A_RBT_B);
190 node = gparent;
191 parent = a_rbt_parent(node);
192 a_rbt_set_parent_color(node, parent, A_RBT_R);
193 continue;
195 tmp = parent->left;
196 if (node == tmp)
198 /* Case 2 - right rotate at parent */
199 tmp = node->right;
200 parent->left = tmp;
201 node->right = parent;
202 if (tmp) { a_rbt_set_parent_color(tmp, parent, A_RBT_B); }
203 a_rbt_set_parent_color(parent, node, A_RBT_R);
204 parent = node;
205 tmp = node->left;
207 /* Case 3 - left rotate at gparent */
208 gparent->right = tmp; /* == parent->left */
209 parent->left = gparent;
210 if (tmp) { a_rbt_set_parent_color(tmp, gparent, A_RBT_B); }
211 a_rbt_set_parents(root, gparent, parent, A_RBT_R);
212 break;
217 static A_INLINE void a_rbt_remove_adjust(a_rbt *root, a_rbt_node *parent)
219 for (a_rbt_node *node = A_NULL, *sibling, *tmp1, *tmp2;;)
222 Loop invariants:
223 - node is black (or null on first iteration)
224 - node is not the root (parent is not null)
225 - All leaf paths going through parent and node have a
226 black node count that is 1 lower than other leaf paths.
228 if (node != parent->right) /* node == parent->left */
230 sibling = parent->right;
231 if (a_rbt_color(sibling) == A_RBT_R)
234 Case 1 - left rotate at parent
237 / \ / \
238 N s --> p Sr
239 / \ / \
240 Sl Sr N Sl
242 tmp1 = sibling->left;
243 parent->right = tmp1;
244 sibling->left = parent;
245 a_rbt_set_parent_color(tmp1, parent, A_RBT_B);
246 a_rbt_set_parents(root, parent, sibling, A_RBT_R);
247 sibling = tmp1;
249 tmp1 = sibling->right;
250 if (!tmp1 || a_rbt_color(tmp1) == A_RBT_B)
252 tmp2 = sibling->left;
253 if (!tmp2 || a_rbt_color(tmp2) == A_RBT_B)
256 Case 2 - sibling color flip (p could be either color here)
258 (p) (p)
259 / \ / \
260 N S --> N s
261 / \ / \
262 Sl Sr Sl Sr
264 This leaves us violating 5) which can be fixed by flipping p to black if it was red,
265 or by recursing at p. p is red when coming from Case 1.
267 a_rbt_set_parent_color(sibling, parent, A_RBT_R);
268 if (a_rbt_color(parent) == A_RBT_R)
270 a_rbt_set_black(parent);
272 else
274 node = parent;
275 parent = a_rbt_parent(node);
276 if (parent) { continue; }
278 break;
281 Case 3 - right rotate at sibling (p could be either color here)
283 (p) (p)
284 / \ / \
285 N S --> N sl
286 / \ \
287 sl Sr S
291 Note: p might be red, and then both p and sl are red after rotation(which breaks property 4).
292 This is fixed in Case 4 (in a_rbt_set_parents() which set sl the color of p and set A_RBT_B)
294 (p) (sl)
295 / \ / \
296 N sl --> P S
297 \ / \
298 S N Sr
302 tmp1 = tmp2->right;
303 sibling->left = tmp1;
304 tmp2->right = sibling;
305 parent->right = tmp2;
306 if (tmp1) { a_rbt_set_parent_color(tmp1, sibling, A_RBT_B); }
307 tmp1 = sibling;
308 sibling = tmp2;
311 Case 4 - left rotate at parent + color flips (p and sl could be either color here.
312 After rotation, p becomes black, s acquires p's color, and sl keeps its color)
314 (p) (s)
315 / \ / \
316 N S --> P Sr
317 / \ / \
318 (sl) sr N (sl)
320 tmp2 = sibling->left;
321 parent->right = tmp2;
322 sibling->left = parent;
323 a_rbt_set_parent_color(tmp1, sibling, A_RBT_B);
324 if (tmp2) { a_rbt_set_parent(tmp2, parent); }
325 a_rbt_set_parents(root, parent, sibling, A_RBT_B);
326 break;
328 else if (parent->left)
330 sibling = parent->left;
331 if (a_rbt_color(sibling) == A_RBT_R)
333 /* Case 1 - right rotate at parent */
334 tmp1 = sibling->right;
335 parent->left = tmp1;
336 sibling->right = parent;
337 a_rbt_set_parent_color(tmp1, parent, A_RBT_B);
338 a_rbt_set_parents(root, parent, sibling, A_RBT_R);
339 sibling = tmp1;
341 tmp1 = sibling->left;
342 if (!tmp1 || a_rbt_color(tmp1) == A_RBT_B)
344 tmp2 = sibling->right;
345 if (!tmp2 || a_rbt_color(tmp2) == A_RBT_B)
347 /* Case 2 - sibling color flip */
348 a_rbt_set_parent_color(sibling, parent, A_RBT_R);
349 if (a_rbt_color(parent) == A_RBT_R)
351 a_rbt_set_black(parent);
353 else
355 node = parent;
356 parent = a_rbt_parent(node);
357 if (parent) { continue; }
359 break;
361 /* Case 3 - left rotate at sibling */
362 tmp1 = tmp2->left;
363 sibling->right = tmp1;
364 tmp2->left = sibling;
365 parent->left = tmp2;
366 if (tmp1) { a_rbt_set_parent_color(tmp1, sibling, A_RBT_B); }
367 tmp1 = sibling;
368 sibling = tmp2;
370 /* Case 4 - right rotate at parent + color flips */
371 tmp2 = sibling->right;
372 parent->left = tmp2;
373 sibling->right = parent;
374 a_rbt_set_parent_color(tmp1, sibling, A_RBT_B);
375 if (tmp2) { a_rbt_set_parent(tmp2, parent); }
376 a_rbt_set_parents(root, parent, sibling, A_RBT_B);
377 break;
379 else { break; }
383 void a_rbt_remove(a_rbt *root, a_rbt_node *node)
385 a_rbt_node *child = node->right;
386 a_rbt_node *tmp = node->left;
387 a_rbt_node *parent, *adjust;
388 #if defined(A_SIZE_POINTER) && (A_SIZE_POINTER + 0 > 1)
389 a_uptr pc;
390 #else /* !A_SIZE_POINTER */
391 unsigned int color;
392 #endif /* A_SIZE_POINTER */
394 if (!tmp)
397 Case 1: node to erase has no more than 1 child (easy!)
398 Note that if there is one child it must be red due to 5) and node must be black due to 4).
399 We adjust colors locally so as to bypass a_rbt_remove_adjust() later on.
401 #if defined(A_SIZE_POINTER) && (A_SIZE_POINTER + 0 > 1)
402 pc = node->parent_;
403 #else /* !A_SIZE_POINTER */
404 color = node->color;
405 #endif /* A_SIZE_POINTER */
406 parent = a_rbt_parent(node);
407 a_rbt_new_child(root, parent, node, child);
408 if (child)
410 #if defined(A_SIZE_POINTER) && (A_SIZE_POINTER + 0 > 1)
411 child->parent_ = pc;
412 #else /* !A_SIZE_POINTER */
413 child->parent = parent;
414 child->color = color;
415 #endif /* A_SIZE_POINTER */
416 adjust = A_NULL;
418 else
420 #if defined(A_SIZE_POINTER) && (A_SIZE_POINTER + 0 > 1)
421 adjust = (pc & 1) == A_RBT_B ? parent : A_NULL;
422 #else /* !A_SIZE_POINTER */
423 adjust = color == A_RBT_B ? parent : A_NULL;
424 #endif /* A_SIZE_POINTER */
427 else if (!child)
429 /* Still case 1, but this time the child is node->left */
430 #if defined(A_SIZE_POINTER) && (A_SIZE_POINTER + 0 > 1)
431 tmp->parent_ = node->parent_;
432 #else /* !A_SIZE_POINTER */
433 tmp->parent = node->parent;
434 tmp->color = node->color;
435 #endif /* A_SIZE_POINTER */
436 parent = a_rbt_parent(node);
437 a_rbt_new_child(root, parent, node, tmp);
438 adjust = A_NULL;
440 else
442 a_rbt_node *successor = child, *child2;
444 tmp = child->left;
445 if (!tmp)
448 Case 2: node's successor is its right child
450 (n) (s)
451 / \ / \
452 (x) (s) -> (x) (c)
456 parent = successor;
457 child2 = successor->right;
459 else
462 Case 3: node's successor is leftmost under node's right child subtree
464 (n) (s)
465 / \ / \
466 (x) (y) -> (x) (y)
468 (p) (p)
470 (s) (c)
474 do {
475 parent = successor;
476 successor = tmp;
477 tmp = tmp->left;
478 } while (tmp);
479 child2 = successor->right;
480 parent->left = child2;
481 successor->right = child;
482 a_rbt_set_parent(child, successor);
485 tmp = node->left;
486 successor->left = tmp;
487 a_rbt_set_parent(tmp, successor);
489 #if defined(A_SIZE_POINTER) && (A_SIZE_POINTER + 0 > 1)
490 pc = node->parent_;
491 #else /* !A_SIZE_POINTER */
492 color = node->color;
493 #endif /* A_SIZE_POINTER */
494 tmp = a_rbt_parent(node);
495 a_rbt_new_child(root, node, successor, tmp);
497 if (child2)
499 a_rbt_set_parent_color(child2, parent, A_RBT_B);
500 adjust = A_NULL;
502 else
504 adjust = a_rbt_color(successor) == A_RBT_B ? parent : A_NULL;
506 #if defined(A_SIZE_POINTER) && (A_SIZE_POINTER + 0 > 1)
507 successor->parent_ = pc;
508 #else /* !A_SIZE_POINTER */
509 successor->parent = tmp;
510 successor->color = color;
511 #endif /* A_SIZE_POINTER */
514 if (adjust) { a_rbt_remove_adjust(root, adjust); }
517 a_rbt_node *a_rbt_insert(a_rbt *root, a_rbt_node *node, int (*cmp)(void const *, void const *))
519 a_rbt_node *parent = root->node;
520 a_rbt_node **link = &root->node;
521 while (*link)
523 parent = *link;
524 int const res = cmp(node, parent);
525 if (res < 0)
527 link = &parent->left;
529 else if (res > 0)
531 link = &parent->right;
533 else { return parent; }
535 *link = a_rbt_init(node, parent);
536 a_rbt_insert_adjust(root, node);
537 return A_NULL;
540 a_rbt_node *a_rbt_search(a_rbt const *root, void const *ctx, int (*cmp)(void const *, void const *))
542 for (a_rbt_node *cur = root->node; cur;)
544 int const res = cmp(ctx, cur);
545 if (res < 0)
547 cur = cur->left;
549 else if (res > 0)
551 cur = cur->right;
553 else { return cur; }
555 return A_NULL;
558 a_rbt_node *a_rbt_head(a_rbt const *root)
560 a_rbt_node *node = root->node;
561 if (node)
563 while (node->left) { node = node->left; }
565 return node;
568 a_rbt_node *a_rbt_tail(a_rbt const *root)
570 a_rbt_node *node = root->node;
571 if (node)
573 while (node->right) { node = node->right; }
575 return node;
578 a_rbt_node *a_rbt_next(a_rbt_node *node)
584 / \ / \
585 A C E G
587 if (!node) { return node; }
588 if (node->right) /* D -> F -> E */
590 for (node = node->right; node->left; node = node->left)
594 else /* C -> B -> D */
596 a_rbt_node *last = node;
597 for (node = a_rbt_parent(node); node && node->left != last;)
599 last = node;
600 node = a_rbt_parent(node);
603 return node;
606 a_rbt_node *a_rbt_prev(a_rbt_node *node)
608 if (!node) { return node; }
609 if (node->left)
611 for (node = node->left; node->right; node = node->right) {}
613 else
615 a_rbt_node *last = node;
616 for (node = a_rbt_parent(node); node && node->right != last;)
618 last = node;
619 node = a_rbt_parent(node);
622 return node;
625 a_rbt_node *a_rbt_pre_next(a_rbt_node *node)
631 / \ / \
632 A C E G
634 a_rbt_node *last = node;
635 if (!node) { return node; }
636 if (node->left) { return node->left; }
637 if (node->right) { return node->right; }
638 for (node = a_rbt_parent(node); node; node = a_rbt_parent(node))
640 if (node->right && node->right != last)
642 node = node->right; /* A -> B -> C */
643 break;
645 last = node; /* C -> B -> D -> F */
647 return node;
650 a_rbt_node *a_rbt_pre_prev(a_rbt_node *node)
652 a_rbt_node *last = node;
653 if (!node) { return node; }
654 if (node->right) { return node->right; }
655 if (node->left) { return node->left; }
656 for (node = a_rbt_parent(node); node; node = a_rbt_parent(node))
658 if (node->left && node->left != last)
660 node = node->left;
661 break;
663 last = node;
665 return node;
668 #define A_RBT_POST(head, tail) \
669 for (;;) \
671 if (node->head) \
673 node = node->head; \
675 else if (node->tail) \
677 node = node->tail; \
679 else { break; } \
681 (void)0
683 a_rbt_node *a_rbt_post_head(a_rbt const *root)
685 a_rbt_node *node = root->node;
686 if (node) { A_RBT_POST(left, right); }
687 return node;
690 a_rbt_node *a_rbt_post_tail(a_rbt const *root)
692 a_rbt_node *node = root->node;
693 if (node) { A_RBT_POST(right, left); }
694 return node;
697 a_rbt_node *a_rbt_post_next(a_rbt_node *node)
703 / \ / \
704 A C E G
706 a_rbt_node *last = node;
707 if (!node) { return node; }
708 node = a_rbt_parent(node);
709 if (node && node->right && node->right != last)
711 node = node->right; /* B -> D -> F -> E */
712 A_RBT_POST(left, right); /* A -> B -> C */
713 } /* C -> B */
714 return node;
717 a_rbt_node *a_rbt_post_prev(a_rbt_node *node)
719 a_rbt_node *last = node;
720 if (!node) { return node; }
721 node = a_rbt_parent(node);
722 if (node && node->left && node->left != last)
724 node = node->left;
725 A_RBT_POST(right, left);
727 return node;
730 a_rbt_node *a_rbt_tear(a_rbt *root, a_rbt_node **next)
732 a_rbt_node *node = *next;
733 if (!node)
735 node = root->node;
736 if (!node) { return node; }
738 A_RBT_POST(left, right);
739 *next = a_rbt_parent(node);
740 a_rbt_new_child(root, *next, node, A_NULL);
741 return node;