Advertisement
Guest User

Untitled

a guest
Nov 4th, 2015
143
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 3.77 KB | None | 0 0
  1. def rot_left(p, li):
  2.     ri = 1 - li
  3.     l = p.lr[li]
  4.     r = p.lr[ri]
  5.     new_p = p.Clone()
  6.     new_r = r.Clone()
  7.  
  8.     new_p.lr[li] = l
  9.     new_p.lr[ri] = r.lr[li]
  10.  
  11.     new_r.lr[li] = new_p
  12.     new_r.lr[ri] = r.lr[ri]
  13.  
  14.     return new_p, new_r
  15.  
  16. def rot_right(p, li):
  17.     return rot_left(p, 1 - li)
  18.  
  19. def is_black(n):
  20.     return n == None or n.color == 'BLACK'
  21.  
  22. def is_red(n):
  23.     return not is_black(n)
  24.  
  25.  
  26. def remove_black_node_with_one_black_child(n, parents):
  27.     # cases numbered according to 2015/Oct/03 wiki RB-tree article
  28.     # https://en.wikipedia.org/w/index.php?title=Red%E2%80%93black_tree&oldid=686983620#Removal
  29.     while 1:
  30.         if not parents:
  31.             return n
  32.         p, li = parents.pop()
  33.         ri = 1 - li
  34.         s = p.lr[ri] # sibling
  35.         sl = s.lr[li]
  36.         sr = s.lr[ri]
  37.         if s.color == 'RED':
  38.             # wiki case 2
  39.             new_p, new_s = rot_left(p, li)
  40.             new_s.color = 'BLACK'
  41.             new_p.lr[li] = n
  42.             new_p.color = 'RED'
  43.             parents.append((new_s, li))
  44.             parents.append((new_p, li))
  45.         elif p.color == s.color and is_black(sl) and is_black(sr):
  46.             # wiki case 3
  47.             new_s = s.Clone()
  48.             new_s.color = 'RED'
  49.             new_p = clone_and_replace(li, n, new_s, p)
  50.             n = new_p
  51.         elif p.color == 'RED' and is_black(s) and is_black(sl) and is_black(sr):
  52.             # wiki case 4
  53.             new_s = s.Clone()
  54.             new_s.color = 'RED'
  55.             new_p = clone_and_replace(li, n, new_s, p, 'BLACK')
  56.             return new_p
  57.         else:
  58.             assert s.color == 'BLACK'
  59.             if is_red(sl) and is_black(sr):
  60.                 assert sl and is_black(sl.lr[li]) and is_black(sl.lr[ri])
  61.                 # wiki case 5
  62.                 new_s, new_sl = rot_right(s, li)
  63.                 assert is_black(new_sl.lr[li])
  64.                 new_p = clone_and_replace(li, p.lr[li], new_sl, p)
  65.                 new_s.color = 'RED'
  66.                 new_sl.color = 'BLACK'
  67.                 assert is_black(new_sl.lr[li])
  68.                 p, s = new_p, new_sl
  69.                 sl, sr = s.lr[li], s.lr[ri]
  70.             assert is_red(sr)
  71.             # wiki case 6
  72.             new_p, new_s = rot_left(p, li)
  73.             new_p.color, new_s.color = new_s.color, new_p.color
  74.             new_sr = sr.Clone()
  75.             new_sr.color = 'BLACK'
  76.             new_s.lr[ri] = new_sr
  77.             new_p.lr[li] = n
  78.             return new_s
  79.     return n
  80.  
  81. def recreate_path_to_new_node(parents, node):
  82.     assert node == None or isinstance(node, TreeNode)
  83.     result = node
  84.     while parents:
  85.         p, ichild = parents.pop()
  86.         result = clone_and_replace(ichild, result, p.lr[1-ichild], p)
  87.     return result
  88.  
  89.  
  90. def remove_node(node, parents):
  91.     if node.lr[0] and node.lr[1]:
  92.         new_n = node.Clone()
  93.         parents.append((new_n, 0))
  94.         n = new_n.lr[0]
  95.         while n.lr[1]:
  96.             parents.append((n, 1))
  97.             n = n.lr[1]
  98.         new_n.key = n.key
  99.         new_n.value = n.value
  100.         node_to_remove = n
  101.         ichild = 0
  102.     else:
  103.         node_to_remove = node
  104.         ichild = int(node.lr[0] == None)
  105.  
  106.     child = node_to_remove.lr[ichild]
  107.     # simple cases with red node - just delete/recolor red node
  108.     if node_to_remove.color == 'RED':
  109.         return recreate_path_to_new_node(parents, child)
  110.     elif node_to_remove.color == 'BLACK':
  111.         if child and child.color == 'RED':
  112.             new_child = child.Clone()
  113.             new_child.color = 'BLACK'
  114.             return recreate_path_to_new_node(parents, new_child)
  115.         else:
  116.             remove_done_node = remove_black_node_with_one_black_child(child, parents)
  117.             return recreate_path_to_new_node(parents, remove_done_node)
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement