Red black tree sort/Wren
This is based, more or less, on the code here which I thought was reasonably intelligible.
Although red black trees don't necessarily need to disallow duplicate keys, I've adjusted the code to optionally disallow them. Not that it matters here as the method of constructing the random keys ensures that there will be no duplicates in any case.
import "random" for Random
/* Represents a node in the red-black tree. */
class Node {
// constructs a new node
construct new(data, parent, left, right, color) {
_data = data
_parent = parent
_left = left
_right = right
_color = color // 0 black, 1 red
}
// constructs a new 'sentinel' node
static new() { new(null, null, null, null, 0) }
data { _data }
parent { _parent }
left { _left }
right { _right }
color { _color }
data=(d) { _data = d }
parent=(p) { _parent = p }
left=(l) { _left = l }
right=(r) { _right = r }
color=(c) { _color = c }
}
/* Represents a red-black tree data structure. */
class RBTree {
// constructs a new red-black tree specifying whether duplicate keys are allowed or not
construct new(allowDups) {
_tnull = Node.new()
_root = _tnull
_allowDups = allowDups
}
// convenience version of the constructor which always allows duplicate keys
static new() { new(true) }
// returns the root node
root { _root }
/* private helper methods for basic operations */
preOrderHelper_(node) {
if (node != _tnull) {
System.print(node.data + " ")
preOrderHelper_(node.left)
preOrderHelper_(node.right)
}
}
inOrderHelper_(node) {
if (node != _tnull) {
inOrderHelper_(node.left)
System.print(node.data + " ")
inOrderHelper_(node.right)
}
}
postOrderHelper_(node) {
if (node != _tnull) {
postOrderHelper_(node.left)
postOrderHelper_(node.right)
System.print(node.data + " ")
}
}
searchTreeHelper_(node, key) {
if (node == _tnull || key == node.data) return node
if (key < node.data) return searchTreeHelper_(node.left, key)
return searchTreeHelper_(node.right, key)
}
fixDelete_(x) {
var s
while (x != _root && x.color == 0) {
if (x == x.parent.left) {
s = x.parent.right
if (s.color == 1) {
s.color = 0
x.parent.color = 1
leftRotate(x.parent)
s = x.parent.right
}
if (s.left.color == 0 && s.right.color == 0) {
s.color = 1
x = x.parent
} else {
if (s.right.color == 0) {
s.left.color = 0
s.color = 1
rightRotate(s)
s = x.parent.right
}
s.color = x.parent.color
x.parent.color = 0
s.right.color = 0
leftRotate(x.parent)
x = _root
}
} else {
s = x.parent.left
if (s.color == 1) {
s.color = 0
x.parent.color = 1
rightRotate(x.parent)
s = x.parent.left
}
if (s.left.color == 0 && s.right.color == 0) {
s.color = 1
x = x.parent
} else {
if (s.left.color == 0) {
s.right.color = 0
s.color = 1
leftRotate(s)
s = x.parent.left
}
s.color = x.parent.color
x.parent.color = 0
s.left.color = 0
rightRotate(x.parent)
x = _root
}
}
}
x.color = 0
}
rbTransplant_(u, v) {
if (!u.parent) {
_root = v
} else if (u == u.parent.left) {
u.parent.left = v
} else {
u.parent.right = v
}
v.parent = u.parent
}
deleteNodeHelper_(node, key) {
// find the node containing key
var z = _tnull
var x
var y
while (node != _tnull) {
if (node.data == key) z = node
if (node.data <= key) {
node = node.right
} else {
node = node.left
}
}
if (z == _tnull) {
System.print("Couldn't find key in the tree.")
return
}
y = z
var yOriginalColor = y.color
if (z.left == _tnull) {
x = z.right
rbTransplant_(z, z.right)
} else if (z.right == _tnull) {
x = z.left
rbTransplant_(z, z.left)
} else {
y = minimum(z.right)
yOriginalColor = y.color
x = y.right
if (y.parent == z) {
x.parent = y
} else {
rbTransplant_(y, y.right)
y.right = z.right
y.right.parent = y
}
rbTransplant_(z, y)
y.left = z.left
y.left.parent = y
y.color = z.color
}
if (yOriginalColor == 0) fixDelete_(x)
}
fixInsert_(k) {
var u
while (k.parent.color == 1) {
if (k.parent == k.parent.parent.right) {
u = k.parent.parent.left // uncle
if (u.color == 1) {
u.color = 0
k.parent.color = 0
k.parent.parent.color = 1
k = k.parent.parent
} else {
if (k == k.parent.left) {
k = k.parent
rightRotate(k)
}
k.parent.color = 0
k.parent.parent.color = 1
leftRotate(k.parent.parent)
}
} else {
u = k.parent.parent.right // uncle
if (u.color == 1) {
u.color = 0
k.parent.color = 0
k.parent.parent.color = 1
k = k.parent.parent
} else {
if (k == k.parent.right) {
k = k.parent
leftRotate(k)
}
k.parent.color = 0
k.parent.parent.color = 1
rightRotate(k.parent.parent)
}
}
if (k == _root) break
}
_root.color = 0
}
printHelper_(root, indent, last) {
if (root != _tnull) {
System.write(indent)
if (last) {
System.write("R----")
indent = indent + " "
} else {
System.write("L----")
indent = indent + "| "
}
var sColor = (root.color == 1) ? "RED" :"BLACK"
System.print(root.data.toString + "(" + sColor + ")")
printHelper_(root.left, indent, false)
printHelper_(root.right, indent, true)
}
}
/* public methods */
// check if a node with a given key already exists in the tree
hasKey(key) {
var node = searchTreeHelper_(_root, key)
return node.data == key
}
// pre-order traversal
// node.left subtree.right subtree
preOrder() { preOrderHelper_(_root) }
// in-order traversal
// left subtree . node . right subtree
inOrder() { inOrderHelper_(_root) }
// post-order traversal
// left subtree . right Subtree . node
postOrder() { postOrderHelper_(_root) }
// search the tree for the key k
// and return the corresponding node
searchTree(k) { searchTreeHelper_(_root, k) }
// find the node with the minimum key
minimum(node) {
while (node.left != _tnull) node = node.left
return node
}
// find the node with the maximum key
maximum(node) {
while (node.right != _tnull) node = node.right
return node
}
// find the successor of a given node
successor(x) {
// if the right subtree is not _tnull
// the successor is the leftmost node in the
// right subtree
if (x.right != _tnull) return minimum(x.right)
// else it is the lowest ancestor of x whose
// right child is also an ancestor of x.
var y = x.parent
while (y != _tnull && x == y.right) {
x = y
y = y.parent
}
return y
}
// find the predecessor of a given node
predecessor(x) {
// if the left subtree is not _tnull
// the predecessor is the rightmost node in the
// left subtree
if (x.left != _tnull) return maximum(x.left)
// else it is the lowest ancestor of x whose
// left child is also an ancestor of x.
var y = x.parent
while (y != _tnull && x == y.left) {
x = y
y = y.parent
}
}
// rotate left at node x
leftRotate(x) {
var y = x.right
x.right = y.left
if (y.left != _tnull) y.left.parent = x
y.parent = x.parent
if (!x.parent) {
_root = y
} else if (x == x.parent.left) {
x.parent.left = y
} else {
x.parent.right = y
}
y.left = x
x.parent = y
}
// rotate right at node x
rightRotate(x) {
var y = x.left
x.left = y.right
if (y.right != _tnull) y.right.parent = x
y.parent = x.parent
if (!x.parent) {
_root = y
} else if (x == x.parent.right) {
x.parent.right = y
} else {
x.parent.left = y
}
y.right = x
x.parent = y
}
// insert the key to the tree in its appropriate position
// and fix the tree
insert(key) {
// ordinary binary search insertion
var y = null
var x = _root
while (x != _tnull) {
y = x
if (key < x.data) {
x = x.left
} else {
x = x.right
}
}
// if duplicates aren't allowed and the key has already been used, simply return
if (!_allowDups && y && y.data == key) return
// new node must be red
var node = Node.new(key, null, _tnull, _tnull, 1)
// y is parent of x
node.parent = y
if (!y) {
_root = node
} else if (node.data < y.data) {
y.left = node
} else {
y.right = node
}
// if new node is a root node, simply return
if (!node.parent) {
node.color = 0
return
}
// if the grandparent is null, simply return
if (!node.parent.parent) return
// fix the tree
fixInsert_(node)
}
// delete the node from the tree
deleteNode(data) {
deleteNodeHelper_(_root, data)
}
// print the tree structure on the terminal
prettyPrint() {
printHelper_(_root, "", true)
}
}
var bst = RBTree.new(false) // disallow duplicates
var rand = Random.new()
// permit random keys from 1 to 99 say
var candidates = (1..99).toList
// shuffle them and insert the first 30 in the RBT
rand.shuffle(candidates)
var insertions = candidates[0..29]
for (k in insertions) bst.insert(k)
System.print("State of the tree after inserting the 30 keys:")
System.print(insertions)
System.print()
bst.prettyPrint()
// now delete 15 of them at random and redisplay the tree
rand.shuffle(insertions)
var deletions = insertions[0..14]
for (k in deletions) bst.deleteNode(k)
System.print("\nState of the tree after deleting the 15 keys:")
System.print(deletions)
System.print()
bst.prettyPrint()
- Output:
Sample output.
State of the tree after inserting the 30 keys:
[85, 57, 51, 50, 86, 39, 95, 49, 44, 48, 21, 83, 11, 59, 88, 66, 4, 40, 24, 82, 63, 22, 37, 32, 91, 74, 28, 75, 62, 81]
R----50(BLACK)
L----39(RED)
| L----21(BLACK)
| | L----11(BLACK)
| | | L----4(RED)
| | R----24(RED)
| | L----22(BLACK)
| | R----32(BLACK)
| | L----28(RED)
| | R----37(RED)
| R----44(BLACK)
| L----40(BLACK)
| R----49(BLACK)
| L----48(RED)
R----83(RED)
L----66(BLACK)
| L----57(RED)
| | L----51(BLACK)
| | R----62(BLACK)
| | L----59(RED)
| | R----63(RED)
| R----75(RED)
| L----74(BLACK)
| R----82(BLACK)
| L----81(RED)
R----86(BLACK)
L----85(BLACK)
R----91(BLACK)
L----88(RED)
R----95(RED)
State of the tree after deleting the 15 keys:
[32, 40, 57, 63, 66, 75, 86, 59, 83, 51, 24, 62, 82, 39, 37]
R----50(BLACK)
L----28(RED)
| L----21(BLACK)
| | L----11(BLACK)
| | | L----4(RED)
| | R----22(BLACK)
| R----48(BLACK)
| L----44(BLACK)
| R----49(BLACK)
R----85(BLACK)
L----81(BLACK)
| L----74(RED)
R----91(RED)
L----88(BLACK)
R----95(BLACK)