sidef/scripts/RosettaCode/avl_tree.sf

#!/usr/bin/ruby

#
## https://rosettacode.org/wiki/AVL_tree
#

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() }

        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() }
        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"
10.times { |i| tree.insert(i) }
print "Printing balance: "
tree.printBalance
print "\n"