Exercise 14.1.6

Observe that whenever we reference the $size$ attribute of a node in either OS-SELECT or OS-RANK, we use it only to compute a rank. Accordingly, suppose we store in each node its rank in the subtree of which it is the root. Show how to maintain this information during insertion and deletion. (Remember that those two operations can cause rotations).

First, let's notice that node.rank = node.left.size + 1. Since one is actually a function of the other, we can use that to calculate the new rank. We need to:

• Update the code doing the rotation, where a slight asymmetry gets introduced.
• Update the code that goes down the tree when inserting to change rank accordingly.
• Update the rank of the node that gets moved to the deleted position with the rank of the deleted node.
• Update the code that goes up the tree when deleting to change rank accordingly.

Specifically for rotations, referring to Figure 13.2, notice that:

• On a left rotation, the rank of $y$ decreases by the rank of $x$ ($\alpha$ gets removed from $y$'s left subtree)
• On a right rotation, the rank of $y$ increases by the rank of $x$ ($\alpha$ gets added to $y$'s left subtree)

All the changes can be found in the code below.

Python code

from enum import Enum
from collections import deque

class Color(Enum):
RED = 1
BLACK = 2

NIL_KEY = object()

def other(direction):
if direction == 'left':
return 'right'
elif direction == 'right':
return 'left'
else:
assert(False)

class Node:
def __init__(self, color, key, parent, left, right, tree, rank=0):
self.color = color
self.key = key
self.parent = parent
self.left = left
self.right = right
self.tree = tree
self.rank = rank

def sexp(self):
if self.isNil():
return 'NIL'

color = 'R' if self.color == Color.RED else 'B'

return f"{color}({self.key}, r={self.rank}, {self.left}, {self.right})"

__str__ = sexp

def black_height(self):
node = self
height = 0

while node is not nil:
if node.color == Color.BLACK:
height += 1
node = node.parent

return height

def isRed(self):
return self.color == Color.RED

def isBlack(self):
return self.color == Color.BLACK

def isNil(self):
return self.key is NIL_KEY

def isNotNil(self):
return not self.isNil()

def __bool__(self):
return self.isNotNil()

def child(self, direction):
if direction == 'left':
return self.left
elif direction == 'right':
return self.right
else:
assert(False)

def set_child(self, direction, child):
if direction == 'left':
self.left = child
elif direction == 'right':
self.right = child
else:
assert(False)

__getitem__ = child
__setitem__ = set_child

def other(self, direction):
return self.child(other(direction))

def rotate(self, direction):
child = self.other(direction)
self[other(direction)] = child[direction]

if child[direction]:
child[direction].parent = self

child.parent = self.parent

if not self.parent:
self.tree.root = child
elif self is self.parent[direction]:
self.parent[direction] = child
else:
self.parent[other(direction)] = child

child[direction] = self
self.parent = child

if direction == 'left':
child.rank += self.rank
else:
self.rank -= child.rank

def depth(self):
n = 0
node = self

while node:
n += 1
node = node.parent

return n

def left_rotate(self):
self.rotate('left')

def right_rotate(self):
self.rotate('right')

def transplant(self, other):
if not self.parent:
self.tree.root = other
elif self is self.parent.left:
self.parent.left = other
else:
self.parent.right = other
other.parent = self.parent

def set(self, parent=None, left=None, right=None, color=None):
if color:
self.color = color
if left is not None:
self.left = left
if right is not None:
self.right = right
if parent is not None:
self.parent = parent

def minimum(self):
node = self

while node.left:
node = node.left

return node

def select(self, i):
node = self

while node.rank != i:
if node.rank < i:
i -= node.rank
node = node.right
else:
node = node.left

return node

def rank_in_tree(self):
rank = self.rank
node = self

while node.parent:
if node.parent.right is node:
rank += node.parent.rank
node = node.parent

return rank

nil = Node(Color.BLACK, NIL_KEY, None, None, None, None, -1)
nil.parent = nil
nil.left = nil
nil.right = nil

class Tree:
def __init__(self):
self.root = nil

def __str__(self):
return self.root.sexp()

def display(self):
height = max([node.depth() for node in self.nodes()])

w = 7
p = 2

level = 0
nodes = [self.root]
while level < height:
cells = (height - 1) ** 2
print(" " * ((height - level) * 5), (" " * p).join((f"{node.key:>d} / {node.rank:>d}" if node else "  NIL  ") for node in nodes))
next_level = []
for node in nodes:
next_level += [node.left, node.right]
level += 1
nodes = next_level

print(f"-- {height} ------")

def search(self, key):
node = self.root

while node:
if node.key == key:
return node
elif key < node.key:
node = node.left
else:
node = node.right

return None

def nodes(self):
items = deque()

if self.root:
items.append(self.root)

while items:
node = items.popleft()

yield node

if node.left:
items.append(node.left)

if node.right:
items.append(node.right)

def select(self, i):
return self.root.select(i)

def insert(self, key):
new = Node(Color.RED, key, None, None, None, self, 1)
parent = nil
node = self.root
while node:
parent = node
if new.key < node.key:
node.rank += 1
node = node.left
else:
node = node.right

new.parent = parent

if not parent:
self.root = new
elif new.key < parent.key:
parent.left = new
else:
parent.right = new

new.set(left=nil, right=nil, color=Color.RED)

self.insert_fixup(new)

def insert_fixup(self, node):
while node.parent.isRed():
if node.parent is node.parent.parent.left:
direction = 'left'
else:
direction = 'right'

if direction == 'left' or direction == 'right':
uncle = node.parent.parent[other(direction)]
if uncle.isRed():
node.parent.color = Color.BLACK
uncle.color = Color.BLACK
node.parent.parent.color = Color.RED
node = node.parent.parent
else:
if node is node.parent[other(direction)]:
node = node.parent
node.rotate(direction)
node.parent.color = Color.BLACK
node.parent.parent.color = Color.RED
node.parent.parent.rotate(other(direction))

self.root.color = Color.BLACK

def delete(self, key):
def decrease_ancestor_ranks(node):
while node:
if node.parent and node.parent.left is node:
node.parent.rank -= 1
node = node.parent

deleted = self.search(key)
y = deleted
y_original_color = y.color

if not deleted.left:
decrease_ancestor_ranks(deleted)
extra_black = deleted.right
deleted.transplant(deleted.right)
elif not deleted.right:
decrease_ancestor_ranks(deleted)
extra_black = deleted.left
deleted.transplant(deleted.left)
else:
y = deleted.right.minimum()
y_original_color = y.color
extra_black = y.right

decrease_ancestor_ranks(y)

if y.parent is deleted:
extra_black.parent = y
else:
y.transplant(y.right)
y.right = deleted.right
y.right.parent = y

deleted.transplant(y)
y.left = deleted.left
y.left.parent = y
y.color = deleted.color
y.rank = deleted.rank

if y_original_color == Color.BLACK:
self.delete_fixup(extra_black)

def delete_fixup(self, node):
while node is not self.root and node.isBlack():
if node is node.parent.left:
direction = 'left'
else:
direction = 'right'

sibling = node.parent[other(direction)]

if sibling.isRed():
sibling.color = Color.BLACK
node.parent.color = Color.RED
node.parent.rotate(direction)
sibling = node.parent[other(direction)]

if sibling.left.isBlack() and sibling.right.isBlack():
sibling.color = Color.RED
node = node.parent
else:
if sibling[other(direction)].isBlack():
sibling[direction].color = Color.BLACK
sibling.color = Color.RED
sibling.rotate(other(direction))
sibling = node.parent[other(direction)]

sibling.color = node.parent.color
node.parent.color = Color.BLACK
sibling[other(direction)].color = Color.BLACK
sibling.parent.rotate(direction)
node = self.root

node.color = Color.BLACK