Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- def rot_left(p, li):
- ri = 1 - li
- l = p.lr[li]
- r = p.lr[ri]
- new_p = p.Clone()
- new_r = r.Clone()
- new_p.lr[li] = l
- new_p.lr[ri] = r.lr[li]
- new_r.lr[li] = new_p
- new_r.lr[ri] = r.lr[ri]
- return new_p, new_r
- def rot_right(p, li):
- return rot_left(p, 1 - li)
- def is_black(n):
- return n == None or n.color == 'BLACK'
- def is_red(n):
- return not is_black(n)
- def remove_black_node_with_one_black_child(n, parents):
- # cases numbered according to 2015/Oct/03 wiki RB-tree article
- # https://en.wikipedia.org/w/index.php?title=Red%E2%80%93black_tree&oldid=686983620#Removal
- while 1:
- if not parents:
- return n
- p, li = parents.pop()
- ri = 1 - li
- s = p.lr[ri] # sibling
- sl = s.lr[li]
- sr = s.lr[ri]
- if s.color == 'RED':
- # wiki case 2
- new_p, new_s = rot_left(p, li)
- new_s.color = 'BLACK'
- new_p.lr[li] = n
- new_p.color = 'RED'
- parents.append((new_s, li))
- parents.append((new_p, li))
- elif p.color == s.color and is_black(sl) and is_black(sr):
- # wiki case 3
- new_s = s.Clone()
- new_s.color = 'RED'
- new_p = clone_and_replace(li, n, new_s, p)
- n = new_p
- elif p.color == 'RED' and is_black(s) and is_black(sl) and is_black(sr):
- # wiki case 4
- new_s = s.Clone()
- new_s.color = 'RED'
- new_p = clone_and_replace(li, n, new_s, p, 'BLACK')
- return new_p
- else:
- assert s.color == 'BLACK'
- if is_red(sl) and is_black(sr):
- assert sl and is_black(sl.lr[li]) and is_black(sl.lr[ri])
- # wiki case 5
- new_s, new_sl = rot_right(s, li)
- assert is_black(new_sl.lr[li])
- new_p = clone_and_replace(li, p.lr[li], new_sl, p)
- new_s.color = 'RED'
- new_sl.color = 'BLACK'
- assert is_black(new_sl.lr[li])
- p, s = new_p, new_sl
- sl, sr = s.lr[li], s.lr[ri]
- assert is_red(sr)
- # wiki case 6
- new_p, new_s = rot_left(p, li)
- new_p.color, new_s.color = new_s.color, new_p.color
- new_sr = sr.Clone()
- new_sr.color = 'BLACK'
- new_s.lr[ri] = new_sr
- new_p.lr[li] = n
- return new_s
- return n
- def recreate_path_to_new_node(parents, node):
- assert node == None or isinstance(node, TreeNode)
- result = node
- while parents:
- p, ichild = parents.pop()
- result = clone_and_replace(ichild, result, p.lr[1-ichild], p)
- return result
- def remove_node(node, parents):
- if node.lr[0] and node.lr[1]:
- new_n = node.Clone()
- parents.append((new_n, 0))
- n = new_n.lr[0]
- while n.lr[1]:
- parents.append((n, 1))
- n = n.lr[1]
- new_n.key = n.key
- new_n.value = n.value
- node_to_remove = n
- ichild = 0
- else:
- node_to_remove = node
- ichild = int(node.lr[0] == None)
- child = node_to_remove.lr[ichild]
- # simple cases with red node - just delete/recolor red node
- if node_to_remove.color == 'RED':
- return recreate_path_to_new_node(parents, child)
- elif node_to_remove.color == 'BLACK':
- if child and child.color == 'RED':
- new_child = child.Clone()
- new_child.color = 'BLACK'
- return recreate_path_to_new_node(parents, new_child)
- else:
- remove_done_node = remove_black_node_with_one_black_child(child, parents)
- return recreate_path_to_new_node(parents, remove_done_node)
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement