def treap_remove_data(node, data):
"""
:type node: TreapNode
"""
if node is None:
return None, None
elif data < node.data:
node.left, removed = treap_remove_data(node.left, data)
return node, removed
elif data > node.data:
node.right, removed = treap_remove_data(node.right, data)
return node, removed
else:
assert data == node.data
if node.right is None:
return node.left, node
elif node.left is None:
return node.right, node
else:
# move node down until it can be removed directly
if node.left.priority < node.right.priority:
# node.left is potentially larger
node = rotate_right(node)
node.right, removed = treap_remove_data(node.right, data)
else:
node = rotate_left(node)
node.left, removed = treap_remove_data(node.left, data)
return node, removed
def rst_rotate_left(node):
"""
:type node: RSNode
"""
node = rotate_left(node)
rst_fix_size(node.left)
rst_fix_size(node)
return node
def rb_insert_node(cur, node):
"""
:type cur: RBNode
:type node: RBNode
"""
if cur is None:
return node
else:
if node.data < cur.data:
cur.left = rb_insert_node(cur.left, node)
if cur.color == RBNode.BLACK:
if cur.left.color == RBNode.RED:
if rb_color_of(cur.left.right) == RBNode.RED:
# B B
# / \ /
# R * R
# / \ ==> / \
# * R R *
# / \ / \
# * * * *
cur.left = rotate_left(cur.left)
if rb_color_of(cur.left.left) == RBNode.RED:
# B
# / \ R
# R * / \
# / \ ==> B B
# R * / \ |\
# / \ * * * *
# * *
cur = rotate_right(cur)
assert cur.left.color == RBNode.RED
cur.left.color = RBNode.BLACK
else:
cur.right = rb_insert_node(cur.right, node)
if cur.right.color == RBNode.RED:
if rb_color_of(cur.right.left) == RBNode.RED:
cur.right = rotate_right(cur.right)
if rb_color_of(cur.right.right) == RBNode.RED:
cur = rotate_left(cur)
assert cur.right.color == RBNode.RED
cur.right.color = RBNode.BLACK
return cur
def treap_insert_node(node, new_node):
"""
:type node: TreapNode
:type new_node: TreapNode
"""
if node is None:
return new_node
if new_node.data < node.data:
node.left = treap_insert_node(node.left, new_node)
if node.left.priority < node.priority:
node = rotate_right(node)
else:
node.right = treap_insert_node(node.right, new_node)
if node.right.priority < node.priority:
node = rotate_left(node)
return node
def treap_root_from_sorted(sorted_stream):
try:
stack = [TreapNode(next(sorted_stream))]
except StopIteration:
return None
while True:
try:
new_node = TreapNode(next(sorted_stream))
except StopIteration:
return stack[0]
assert stack[-1].right is None
stack[-1].right = new_node
stack.append(new_node)
while len(stack) >= 2 and stack[-1].priority < stack[-2].priority:
ch, p = stack.pop(), stack.pop()
p = rotate_left(p)
if stack:
stack[-1].right = p
stack.append(p)
def rb_remove_fix_right(node):
"""
:type node: RBNode
:rtype: (RBNode, bool)
"""
if rb_color_of(node.right) == RBNode.RED:
node.right.color = RBNode.BLACK
return node, False
if node.left.color == RBNode.RED:
node = rotate_right(node)
node.color = RBNode.BLACK
node.right.color = RBNode.RED
node.right, need_fix = rb_remove_fix_right(node.right)
if need_fix:
node, need_fix = rb_remove_fix_right(node)
assert not need_fix
return node, False
elif rb_color_of(node.left.right) == rb_color_of(
node.left.left) == RBNode.BLACK:
node.left.color = RBNode.RED
return node, True
elif rb_color_of(node.left.right) != rb_color_of(
node.left.left) == RBNode.BLACK:
node.left = rotate_left(node.left)
node.left.color = RBNode.BLACK
node.left.left.color = RBNode.RED
node, need_fix = rb_remove_fix_right(node)
assert not need_fix
return node, False
else:
assert node.left.left.color == RBNode.RED
node = rotate_right(node)
node.color = node.right.color
node.right.color = node.left.color = RBNode.BLACK
return node, False
def rb_remove_fix_left(node):
"""
:type node: RBNode
:rtype: (RBNode, bool)
"""
if rb_color_of(node.left) == RBNode.RED:
node.left.color = RBNode.BLACK
return node, False
if node.right.color == RBNode.RED:
# B B
# / \ / \
# * B R R B
# /| / \ ==> / \ |\
# 1 1 B B * B B 2 2
# /| / \ /| /|
# 2 2 2 2 1 1 2 2
node = rotate_left(node)
node.color = RBNode.BLACK
node.left.color = RBNode.RED
node.left, need_fix = rb_remove_fix_left(node.left)
if need_fix:
node, need_fix = rb_remove_fix_left(node)
assert not need_fix
return node, False
elif rb_color_of(node.right.left) == rb_color_of(
node.right.right) == RBNode.BLACK:
# ? * ?
# / \ / \
# * B B B R
# /| / \ ==> /| / \
# 1 1 B B 1 1 B B
# /| / \ /| / \
# 1 1 1 1 1 1 1 1
node.right.color = RBNode.RED
return node, True
elif rb_color_of(node.right.left) != rb_color_of(
node.right.right) == RBNode.BLACK:
# ? ? ?
# / \ / \ / \
# * B B B R * B B
# /| / \ ==> /| / \ ==> /| / \
# 1 1 R B 1 1 2 B 1 1 2 R
# /| / \ / \ / \
# 2 2 1 1 2 B 2 B
# / \ / \
# 1 1 1 1
node.right = rotate_right(node.right)
node.right.color = RBNode.BLACK
node.right.right.color = RBNode.RED
node, need_fix = rb_remove_fix_left(node)
assert not need_fix
return node, False
else:
# ? B ?
# / \ / \ / \
# * B B ? R B B
# /| / \ ==> / \ |\ ==> / \ |\
# 1 1 2 R B 2 2 2 B 2 2 2
# / \ / \ / \
# 2 2 1 1 1 1
assert node.right.right.color == RBNode.RED
node = rotate_left(node)
node.color = node.left.color
node.left.color = node.right.color = RBNode.BLACK
return node, False
def avl_rotate_left(node: AVLNode):
assert avl_height_of(node.right) > avl_height_of(node.left)
node = rotate_left(node)
avl_reheight(node.left)
avl_reheight(node)
return node