class AVLtree {
has root = nil
struct Node {
Number key,
Number balance = 0,
Node left = nil,
Node right = nil,
Node parent = nil,
}
method insert(key) {
if (root == nil) {
root = Node(key)
return true
}
var n = root
var parent = nil
loop {
if (n.key == key) {
return false
}
parent = n
var goLeft = (n.key > key)
n = (goLeft ? n.left : n.right)
if (n == nil) {
var tn = Node(key, parent: parent)
if (goLeft) {
parent.left = tn
}
else {
parent.right = tn
}
self.rebalance(parent)
break
}
}
return true
}
method delete_key(delKey) {
if (root == nil) { return nil }
var n = root
var parent = root
var delNode = nil
var child = root
while (child != nil) {
parent = n
n = child
child = (delKey >= n.key ? n.right : n.left)
if (delKey == n.key) {
delNode = n
}
}
if (delNode != nil) {
delNode.key = n.key
child = (n.left != nil ? n.left : n.right)
if (root.key == delKey) {
root = child
}
else {
if (parent.left == n) {
parent.left = child
}
else {
parent.right = child
}
self.rebalance(parent)
}
}
}
method rebalance(n) {
if (n == nil) { return nil }
self.setBalance(n)
given (n.balance) {
when (-2) {
if (self.height(n.left.left) >= self.height(n.left.right)) {
n = self.rotate(n, :right)
}
else {
n = self.rotate_twice(n, :left, :right)
}
}
when (2) {
if (self.height(n.right.right) >= self.height(n.right.left)) {
n = self.rotate(n, :left)
}
else {
n = self.rotate_twice(n, :right, :left)
}
}
}
if (n.parent != nil) {
self.rebalance(n.parent)
}
else {
root = n
}
}
method rotate(a, dir) {
var b = (dir == :left ? a.right : a.left)
b.parent = a.parent
(dir == :left) ? (a.right = b.left)
: (a.left = b.right)
if (a.right != nil) {
a.right.parent = a
}
b.$dir = a
a.parent = b
if (b.parent != nil) {
if (b.parent.right == a) {
b.parent.right = b
}
else {
b.parent.left = b
}
}
self.setBalance(a, b)
return b
}
method rotate_twice(n, dir1, dir2) {
n.left = self.rotate(n.left, dir1)
self.rotate(n, dir2)
}
method height(n) {
if (n == nil) { return -1 }
1 + Math.max(self.height(n.left), self.height(n.right))
}
method setBalance(*nodes) {
nodes.each { |n|
n.balance = (self.height(n.right) - self.height(n.left))
}
}
method printBalance {
self.printBalance(root)
}
method printBalance(n) {
if (n != nil) {
self.printBalance(n.left)
print(n.balance, ' ')
self.printBalance(n.right)
}
}
}
var tree = AVLtree()
say "Inserting values 1 to 10"
{|i| tree.insert(i) } << 1..10
print "Printing balance: "
tree.printBalance
Output:
Inserting values 1 to 10
Printing balance: 0 0 0 1 0 0 0 0 1 0