avl_tree.rb

#!/usr/bin/env ruby

class AVLTree
    class Node
        attr_accessor :key, :data, :height, :left, :right

        def initialize key, data
            self.key    = key
            self.data  = data
            self.height = 1
        end
    end

    def self.height node
        node&.height || 0
    end

    def self.recalc_height(node)
        lh  = AVLTree.height(node.left)
        rh  = AVLTree.height(node.right)
        max = lh > rh ? lh : rh

        node.height = 1 + max
    end

    def self.rotate_right p
        q       = p.left
        p.left  = q.right
        q.right = p

        AVLTree.recalc_height(p)
        AVLTree.recalc_height(q)

        return q
    end

    def self.rotate_left p
        q       = p.right
        p.right = q.left
        q.left  = p

        AVLTree.recalc_height(p)
        AVLTree.recalc_height(q)

        return q
    end

    def self.balance node
        AVLTree.recalc_height node

        if (AVLTree.height(node.left) - AVLTree.height(node.right) == 2) 
            if (AVLTree.height(node.left.right) > AVLTree.height(node.left.left)) 
                node.left = AVLTree.rotate_left(node.left)
            end
            return AVLTree.rotate_right(node)
        elsif (AVLTree.height(node.right) - AVLTree.height(node.left) == 2) 
            if (AVLTree.height(node.right.left) > AVLTree.height(node.right.right)) 
                node.right = AVLTree.rotate_right(node.right)
            end
            return AVLTree.rotate_left(node)
        end
        return node
    end

    def self.search node, key
        return nil unless node
        return AVLTree.search(node.left, key)  if (key < node.key)
        return AVLTree.search(node.right, key) if (key > node.key)
        return node        
    end

    def self.insert node, key, data
        return Node.new(key, data) unless node

        if (key < node.key) 
            node.left = AVLTree.insert(node.left, key, data)
        elsif(key > node.key) 
            node.right = AVLTree.insert(node.right, key, data)
        else 
            node.data = data
        end

        return AVLTree.balance(node)
    end

    def  self.find_min node
        return AVLTree.find_min node.left if node.left
        node
    end

    def  self.remove_min node
        return node.right unless node.left
        node.left = AVLTree.remove_min(node.left)
        AVLTree.balance(node)
    end


    def self.remove_item node, key
        return nil unless node

        if key < node.key
            node.left = AVLTree.remove_item(node.left, key)
        elsif (key > node.key)
            node.right = AVLTree.remove_item(node.right, key)
        else
            left, right = node.left, node.right
            return left unless right

            m       = AVLTree.find_min(right)
            m.right = AVLTree.remove_min(right)
            m.left  = left

            return AVLTree.balance(m)
        end

        return AVLTree.balance(node)
    end

end


# How to use it.
root = nil
node = AVLTree.insert root, 1, "foo"
node = AVLTree.insert node, 2, "bar"
node = AVLTree.insert node, 3, "baz"
node = AVLTree.insert node, 4, "bus"

puts AVLTree.search(node, 2).data # -> "bar"
puts AVLTree.search(node, 3).data # -> "baz"