175 lines
3.8 KiB
Plaintext
175 lines
3.8 KiB
Plaintext
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
|