Skip to content

Instantly share code, notes, and snippets.

@nilium
Created June 1, 2010 02:26
Show Gist options
  • Save nilium/420492 to your computer and use it in GitHub Desktop.
Save nilium/420492 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
}
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) }
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) -> Int
init: func(comp: Func (K, K) -> Int) {
comparator = comp
}
get: func (key: K) -> V {
node := root
while (node) {
cmp := comparator(key as K, node key as K)
"get - comparison result: %d" format(cmp) println()
if (cmp == 0)
return node value
else if (cmp < 0)
node = node left
else
node = node right
}
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
}
}
__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
}
}
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 (left, right: 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(5)) println()
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment