Last active
June 9, 2018 20:52
-
-
Save yuga/4362218 to your computer and use it in GitHub Desktop.
CLRS 2ndのRed-Black Tree
あってるかどうか... 番兵のparent, left, right, keyは操作の都合上値置場として活用することがある(日本語版第二版P269)。
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
#!/bin/env python | |
# -*- coding: utf-8 -*- | |
class Color: | |
RED = 1 | |
BLACK = 2 | |
class Element: | |
def __init__(self, key): | |
self.key = key | |
self.color = None | |
self.parent = None | |
self.left = None | |
self.right = None | |
return | |
def __unicode__(self): | |
if self.color == Color.RED: | |
c = u'R' | |
else: | |
c = u'B' | |
#return u'%s:%s (%s) (%s)' % (c, self.key, self.left, self.right) | |
return u'%s:%s' % (c, self.key) | |
def __str__(self): | |
return unicode(self).encode('utf_8') | |
class RedBlackTree: | |
class NilElement: | |
def __init__(self, tree): | |
self._tree = tree | |
self.color = Color.BLACK | |
return | |
#@property | |
#def parent(self): | |
# return None | |
#@parent.getter | |
#def parent(self): | |
# return self._tree.root | |
def __init__(self): | |
self.nil = RedBlackTree.NilElement(self) | |
self.root = self.nil | |
return | |
def left_rotate(self, x): | |
y = x.right | |
x.right = y.left # yの左部分木をxの右部分木にする | |
if y.left != self.nil: | |
y.left.parent = x | |
y.parent = x.parent # xの親をyにリンクする | |
if x.parent == self.nil: | |
self.root = y | |
elif x == x.parent.left: # xが親の左の子だったら | |
x.parent.left = y | |
else: # xが親の右の子だったら | |
x.parent.right = y | |
y.left = x # xをyの左におく | |
x.parent = y | |
return | |
def right_rotate(self, x): | |
y = x.left | |
x.left = y.right # yの右部分木をxの左部分木にする | |
if y.right != self.nil: | |
y.right.parent = x | |
y.parent = x.parent # xの親をyにリンクする | |
if x.parent == self.nil: | |
self.root = y | |
elif x == x.parent.left: # xが親の左の子だったら | |
x.parent.left = y | |
else: # xが親の右の子だったら | |
x.parent.right = y | |
y.right = x # xをyの右におく | |
x.parent = y | |
return | |
def insert(self, z): | |
y = self.nil | |
x = self.root | |
while x != self.nil: # 葉に到達するまで | |
y = x | |
if z.key < x.key: # keyを比較して左右に振りわけ | |
x = x.left | |
else: | |
x = x.right | |
z.parent = y | |
if y == self.nil: # 空の木だった場合 | |
self.root = z | |
elif z.key < y.key: # zがyの左の子になる場合 | |
y.left = z | |
else: # zがyの右の子になる場合 | |
y.right = z | |
z.left = self.nil | |
z.right = self.nil | |
z.color = Color.RED # 新しく葉になるzはとりあえずREDにしておく(あとで調整) | |
self._insert_fixup(z) # 葉から根へ向けて調整する | |
return | |
def _insert_fixup(self, z): | |
while z.parent.color == Color.RED: # zの親の色もREDなら | |
if z.parent == z.parent.parent.left: # zの親がその親の左の子の場合 | |
y = z.parent.parent.right | |
if y.color == Color.RED: # [場合1] zの親の親の右の子がREDの場合 | |
z.parent.color = Color.BLACK # [場合1] zの親をBLACKに | |
y.color = Color.BLACK # [場合1] zの親の親の右の子もBLACKに | |
z.parent.parent.color = Color.RED # [場合1] zの親の親をREDに | |
z = z.parent.parent # [場合1] zの親の親へ移動、新たなzにする | |
else: | |
if z == z.parent.right: # [場合2] zが親の右の子の場合 | |
z = z.parent # [場合2] zの親へ移動、新たなzにする | |
self.left_rotate(z) # [場合2] 新たなzで左回転、旧zが親に新zがその左の子に | |
z.parent.color = Color.BLACK # [場合2,3] 旧zをBLACKに | |
z.parent.parent.color = Color.RED # [場合2,3] 新zの親の親をREDに | |
self.right_rotate(z.parent.parent) # [場合2,3] 新zの親の親で右回転(新zの親の親は旧zの右の子に) | |
else: # zが親の右の子の場合 | |
y = z.parent.parent.left | |
if y.color == Color.RED: | |
z.parent.color = Color.BLACK | |
y.color = Color.BLACK | |
z.parent.parent.color = Color.RED | |
z = z.parent.parent | |
else: | |
if z == z.parent.left: | |
z = z.parent | |
self.right_rotate(z) | |
z.parent.color = Color.BLACK | |
z.parent.parent.color = Color.RED | |
self.left_rotate(z.parent.parent) | |
self.root.color = Color.BLACK | |
return | |
def delete(self, z): | |
if z.left == self.nil or z.right == self.nil: # zに左か右の子がないならば | |
y = z # z自身をyとする | |
else: # zにどちらの子もあるならば | |
y = self.successor(z) # zの次の要素をyとする | |
if y.left != self.nil: # yに左の子があれば(z自身がy) | |
x = y.left # yの左の子をx | |
else: # でなければ | |
x = y.right # yの右の子をx(右の子がない場合もある、そのときはnil) | |
x.parent = y.parent # xの親をyの親に | |
if y.parent == self.nil: # yの親がなければ | |
self.root = x # xを根に | |
else: | |
if y == y.parent.left: # yが左の子なら | |
y.parent.left = x # xをyの親の左の子に | |
else: # でなければ | |
y.parent.right = x # xをyの親の右の子に | |
if y != z: # yがzでなければ(zにどちらの子もあるならば) | |
z.key = y.key # yの内容をzに引き写す | |
if y.color == Color.BLACK: # yがBLACKだったら | |
self._delete_fixup(x) # 調整する | |
return y | |
def _delete_fixup(self, x): | |
while x != self.root and x.color == Color.BLACK: # xが根またはREDになるまで | |
if x == x.parent.left: # xが左の子の場合 | |
w = x.parent.right # wを兄弟(右の子) | |
if w.color == Color.RED: # [場合1] wがREDなら | |
w.color = Color.BLACK # [場合1] wをBLACK | |
x.parent.color = Color.RED # [場合1] xの親をRED | |
self.left_rotate(x.parent) # [場合1] xの親で左回転(wがxの親の親になる) | |
w = x.parent.right # [場合1] xの親の新しい右の子(wの左の子だった)をw | |
if w.left.color == Color.BLACK and w.right.color == Color.BLACK: # [場合2] wの子がどちらもBLACK | |
w.color = Color.RED # [場合2] wをRED | |
x = x.parent # [場合2] xの親へ移動 | |
else: # wの子がどちらかRED | |
if w.right.color == Color.BLACK: # [場合3] wの右の子がBLACKなら | |
w.left.color = Color.BLACK # [場合3] wの左の子をBLACK | |
w.color = Color.RED # [場合3] wをRED | |
self.right_rotate(w) # [場合3] wで右回転、wの左の子がwの親に | |
w = x.parent.right # [場合3] xの親の新しい右の子をw | |
w.color = x.parent.color # [場合4] wの色をxの親と同じ色に | |
x.parent.color = Color.BLACK # [場合4] xの親をBLACK | |
w.right.color = Color.BLACK # [場合4] wの右の子をBLACK | |
self.left_rotate(x.parent) # [場合4] xの親で右回転、xが親に | |
x = self.root # [場合4] 根をx | |
else: # xが右の子の場合 | |
w = x.parent.left | |
if w.color == Color.RED: | |
w.color = Color.BLACK | |
x.parent.color = Color.RED | |
self.right_rotate(x.parent) | |
w = x.parent.left | |
if w.left.color == Color.BLACK and w.right.color == Color.BLACK: | |
w.color = Color.RED | |
x = x.parent | |
else: | |
if w.left.color == Color.BLACK: | |
w.right.color = Color.BLACK | |
w.color = Color.RED | |
self.left_rotate(w) | |
w = w.parent.left | |
w.color = x.parent.color | |
x.parent.color = Color.BLACK | |
w.right.color = Color.BLACK | |
self.right_rotate(x.parent) | |
x = self.root | |
x.color = Color.BLACK | |
return | |
def successor(self, x): | |
if x.right != self.nil: | |
return self.minimum(x.right) | |
y = x.parent | |
while y != self.nil and x == y.right: | |
x = y | |
y = y.parent | |
return y | |
def minimum(self, x): | |
while x.left != self.nil: | |
x = x.left | |
return x | |
def search(self, x, k): | |
while x != self.nil and k != x.key: | |
if k < x.key: | |
x = x.left | |
else: | |
x = x.right | |
return x | |
def pprint(self): | |
def pr(d, e): | |
print u'%s%s' % (u' ' * d * 4, e) | |
return | |
depth = 0 | |
stack = [] | |
e = self.root | |
if e != self.nil: | |
while True: | |
while e != self.nil: | |
pr(depth, e) | |
stack.append((e,depth)) | |
e = e.left | |
depth += 1 | |
pr(depth, u'NIL') | |
if len(stack) == 0: | |
break | |
(e,depth) = stack.pop() | |
e = e.right | |
depth += 1 | |
else: | |
p(len(stack), u'NIL') | |
return |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment