15.3.2 Step 3: Rebalance

When we rebalanced ordinary RB trees, we used the expressions pa[k - 1] and pa[k - 2] to refer to the parent and grandparent, respectively, of the node at which we were rebalancing, and we called that node q, though that wasn't a variable name (see Inserting an RB Node Step 3 - Rebalance). Now that we have parent pointers, we use a real variable q to refer to the node where we're rebalancing.

This means that we could refer to its parent and grandparent as q->prb_parent and q->prb_parent->prb_parent, respectively, but there's a small problem with that. During rebalancing, we will need to move nodes around and modify parent pointers. That means that q->prb_parent and q->prb_parent->prb_parent will be changing under us as we work. This makes writing correct code hard, and reading it even harder. It is much easier to use a pair of new variables to hold q's parent and grandparent.

That's exactly the role that f and g, respectively, play in the code below. If you compare this code to <Step 3: Rebalance after RB insertion 201>, you'll also notice the way that checking that f and g are non-null corresponds to checking that the stack height is at least 3 (see Exercise 6.4.3-1 for an explanation of the reason this is a valid test).

557. <Step 3: Rebalance after PRB insertion 557> =
q = n;
for (;;) 
  {
    struct prb_node *f; /* Parent of q. */
    struct prb_node *g; /* Grandparent of q. */

    f = q->prb_parent;
    if (f == NULL || f->prb_color == PRB_BLACK)
      break;

    g = f->prb_parent;
    if (g == NULL)
      break;

    if (g->prb_link[0] == f)
      { 
        <Left-side rebalancing after PRB insertion 558> 
      }
    else 
      { 
        <Right-side rebalancing after PRB insertion 562> 
      } 
  }
tree->prb_root->prb_color = PRB_BLACK;

This code is included in 555.

After replacing pa[k - 1] by f and pa[k - 2] by g, the cases for PRB rebalancing are distinguished on the same basis as those for RB rebalancing (see <Left-side rebalancing after RB insertion 202>). One addition: cases 2 and 3 need to work with q's great-grandparent, so they stash it into a new variable h.

558. <Left-side rebalancing after PRB insertion 558> =
struct prb_node *y = g->prb_link[1];
if (y != NULL && y->prb_color == PRB_RED) 
  { 
    <Case 1 in left-side PRB insertion rebalancing 559> 
  }
else 
  {
    struct prb_node *h; /* Great-grandparent of q. */

    h = g->prb_parent;
    if (h == NULL)
      h = (struct prb_node *) &tree->prb_root;

    if (f->prb_link[1] == q)
      { 
        <Case 3 in left-side PRB insertion rebalancing 561> 
      }

    <Case 2 in left-side PRB insertion rebalancing 560>
    break;
  }

This code is included in 557.

Case 1: q's uncle is red

In this case, as before, we need only rearrange colors (see rbinscase1). Instead of popping the top two items off the stack, we directly set up q, the next node at which to rebalance, to be the (former) grandparent of the original q.

559. <Case 1 in left-side PRB insertion rebalancing 559> =
f->prb_color = y->prb_color = PRB_BLACK;
g->prb_color = PRB_RED;
q = g;

This code is included in 558.

Case 2: q is the left child of its parent

If q is the left child of its parent, we rotate right at g:

The result satisfies both RB balancing rules. Refer back to the discussion of the same case in ordinary RB trees for more details (see rbinscase2).

560. <Case 2 in left-side PRB insertion rebalancing 560> =
g->prb_color = PRB_RED;
f->prb_color = PRB_BLACK;

g->prb_link[0] = f->prb_link[1];
f->prb_link[1] = g;
h->prb_link[h->prb_link[0] != g] = f;

f->prb_parent = g->prb_parent;
g->prb_parent = f;
if (g->prb_link[0] != NULL)
  g->prb_link[0]->prb_parent = g;

This code is included in 558.

Case 3: q is the right child of its parent

If q is a right child, then we transform it into case 2 by rotating left at f:

Afterward we relabel q as f and treat the result as case 2. There is no need to properly set q itself because case 2 never uses variable q. For more details, refer back to case 3 in ordinary RB trees (see rbinscase3).

561. <Case 3 in left-side PRB insertion rebalancing 561> =
f->prb_link[1] = q->prb_link[0];
q->prb_link[0] = f;
g->prb_link[0] = q;
f->prb_parent = q;
if (f->prb_link[1] != NULL)
  f->prb_link[1]->prb_parent = f;

f = q;

This code is included in 558.