Skip to content

Instantly share code, notes, and snippets.

@yuga
Last active June 9, 2018 20:52
Show Gist options
  • Save yuga/4362218 to your computer and use it in GitHub Desktop.
Save yuga/4362218 to your computer and use it in GitHub Desktop.
CLRS 2ndのRed-Black Tree あってるかどうか... 番兵のparent, left, right, keyは操作の都合上値置場として活用することがある(日本語版第二版P269)。
#!/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