Skip to content

Instantly share code, notes, and snippets.

@komiga
Forked from nilium/rbtree.ooc
Created June 1, 2010 18:27
Show Gist options
  • Save komiga/421286 to your computer and use it in GitHub Desktop.
Save komiga/421286 to your computer and use it in GitHub Desktop.
RBNodeColor: enum {
black = 0
red = 1
}
RBNode: class <K, V> {
color := RBNodeColor black
left, right, parent: This<K, V>
key: K
value: V
key: func() -> K { key }
value: func() -> V { value }
init: func(=key, =value, =parent) {}
grandparent: func -> This<K, V> {
if (parent)
return parent parent
return null
}
uncle: func -> This<K, V> {
gp := grandparent()
if (gp == null)
return null
if (gp left == this)
return gp right
else
return gp left
}
sibling: func -> This<K, V> {
if (parent) {
if (isLeft?())
return parent right
else
return parent left
}
return null
}
isLeft?: func -> Bool { (parent && parent left == this) }
isRight?: func -> Bool { (parent && parent right == this) }
isBlack?: func -> Bool { (color == RBNodeColor black) }
isRed?: func -> Bool { (color == RBNodeColor red) }
hasParent?: func -> Bool { (parent != null) }
isLeaf?: func -> Bool { (!left && !right) }
hasLeft?: func -> Bool { (left != null) }
hasRight?: func -> Bool { (right != null) }
walk: func (fn: Func(RBNode<K,V>)) {
if (left) {
left walk(fn)
}
fn(this)
if (right) {
right walk(fn)
}
}
}
RBTree: class <K, V> {
root: RBNode<K, V> = null
comparator: Func<K> (K, K) -> Int
init: func(comp: Func<K> (K, K) -> Int) {
comparator = comp
}
replaceNode: func (oldn, newn: RBNode<K, V>) {
if (!oldn hasParent?()) {
root = newn
} else {
if (oldn isLeft?())
oldn parent left = newn
else
oldn parent right = newn
}
if (newn)
newn parent = oldn parent
}
contains: func (key: K) -> Bool {
return __findNode(key) ? true : false
}
get: func (key: K) -> V {
node := __findNode(key)
if (node)
return node value
else
return null
}
put: func (key: K, value: V) {
node := RBNode<K, V> new(key, value, null)
parent: RBNode<K, V> = null
follower := root
cmp: Int = 0
while (follower) {
parent = follower
cmp = comparator(key as K, follower key as K)
"put - comparison result: %d" format(cmp) println()
if (cmp < 0) {
follower = follower left
} else if (cmp > 0) {
follower = follower right
} else {
follower value = value
return
}
}
if (parent) {
node parent = parent
cmp := comparator(node key as K, parent key as K)
if (cmp < 0)
parent left = node
else
parent right = node
} else {
root = node
}
node color = RBNodeColor red
while (node) {
// mostly based on the wikipedia page and this page:
// http://www.cs.auckland.ac.nz/software/AlgAnim/red_black.html
if (node parent == null) {
// case 1
node color = RBNodeColor black
return
} else if (node parent isBlack?()) {
// case 2
return
} else {
// case 3
uncle := node uncle()
if (uncle && uncle isRed?()) {
node parent color = RBNodeColor black
uncle color = RBNodeColor black
gpar := node grandparent()
gpar color = RBNodeColor red
node = gpar
continue
} else {
// case 4
gpar := node grandparent()
if (node isRight?() && node parent isLeft?()) {
__rotateLeft(node parent)
node = node left
} else if (node isLeft?() && node parent isRight?()) {
__rotateRight(node parent)
node = node right
}
// case 5
gpar = node grandparent()
node parent color = RBNodeColor black
gpar color = RBNodeColor red
if (node isLeft?() && node parent isLeft?()) {
__rotateRight(gpar)
} else if (node isRight?() && node parent isRight?()) {
__rotateLeft(gpar)
}
}
}
return
}
}
remove: func (key: K) -> Bool {
n := __findNode(key)
if (n) {
if (!n isLeaf?()) {
// Can't use 'n right isLeaf?()' here because the right node might be null
child := n hasLeft?() ? n left : n right
replaceNode(n, child)
if (n isBlack?()) {
if (child isRed?())
child color = RBNodeColor black
else
__deleteNode(child)
}
} else {
replaceNode(n, null)
}
return true
}
return false
}
__deleteNode: func (n: RBNode<K, V>) {
if (n parent != null) {
// case 2
s := n sibling()
if (s isRed?()) {
n parent color = RBNodeColor red
s color = RBNodeColor black
if (n isLeft?())
__rotateLeft(n parent)
else
__rotateRight(n parent)
}
// case 3
if ((n parent isBlack?()) &&
(s isBlack?()) &&
(s left isBlack?()) &&
(s right isBlack?())) {
s color = RBNodeColor red
__deleteNode(n parent)
} else {
// case 4
if ((n parent isRed?()) &&
(s isBlack?()) &&
(s left isBlack?()) &&
(s right isBlack?())) {
s color = RBNodeColor red
n parent color = RBNodeColor black
} else {
// case 5
if (s isBlack?()) {
if ((n isLeft?()) &&
(s right isBlack?()) &&
(s left isRed?())) {
s color = RBNodeColor red
s left color = RBNodeColor black
__rotateRight(s)
} else if ((n isRight?()) &&
(s left isBlack?()) &&
(s right isRed?())) {
s color = RBNodeColor red
s right color = RBNodeColor black
__rotateLeft(s)
}
}
// case 6
s color = n parent color
n parent color = RBNodeColor black
if (n isLeft?()) {
s right color = RBNodeColor black
__rotateLeft(n parent)
} else {
s left color = RBNodeColor black
__rotateRight(n parent)
}
}
}
}
}
__findNode: func (key: K) -> RBNode<K, V> {
node := root
while (node) {
cmp := comparator(key as K, node key as K)
"findNode - comparison result: %d" format(cmp) println()
if (cmp == 0)
return node
else if (cmp < 0)
node = node left
else
node = node right
}
return null
}
__rotateLeft: func (node: RBNode<K, V>) {
right := node right
node right = right left
if (node right)
node right parent = node
right parent = node parent
if (node parent == null) {
root = right
} else {
if (node isLeft?()) {
node parent left = right
} else { // right
node parent right = right
}
}
right left = node
node parent = right
}
__rotateRight: func (node: RBNode<K, V>) {
left := node left
node left = left right
if (node left)
node left parent = node
left parent = node parent
if (node parent == null) {
root = left
} else {
if (node isLeft?()) {
node parent left = left
} else { // right
node parent right = left
}
}
left right = node
node parent = left
}
}
assert: func (test: Bool) {
Exception new("FAILURE") throw()
}
walkFn: func(n: RBNode<Int, Int>) {
color: String = n isBlack?() ? "black" : "red"
if (n isLeft?())
"left: %d %s" format(n key(), color) println()
else if (n isRight?())
"right: %d %s" format(n key(), color) println()
else
"root: %d %s" format(n key(), color) println()
"value: %d" format(n value()) println()
}
compareFn: func<K> (left, right: K) -> Int {
result: Int = (left - right)
"comparing %d and %d - result: %d" format(left, right, result) println()
return result
}
main: func {
woop := RBTree<Int, Int> new(compareFn)
woop put(1, 100)
woop put(2, 200)
woop put(3, 300)
woop put(4, 400)
woop put(5, 500)
woop put(6, 600)
woop put(7, 700)
woop put(8, 800)
woop put(9, 900)
woop put(10, 1000)
woop put(7, 12345)
woop root walk(walkFn)
"%d" format(woop get(6)) println()
woop remove(6)
"value at 6: %d" format(woop get(6)) println()
"contains 6: %d" format(woop contains(6)) println()
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment