Avl Tree in Ruby

avl.rb

# Node
class Node
  attr_accessor :value,
                :height,
                :left,
                :right

  def initialize(value = nil, height = 0)
    @value  = value
    @height = height
  end
end

# AvlTree
class AvlTree
  attr_accessor :root

  def initialize(root = nil)
    @root = root
  end

  def height(node = nil)
    return -1 unless node
    node.height
  end

  def put(value, node = root)
    return @root = _put(value) unless node

    _put(value, node)
  end

  def _put(value, node = nil)
    return Node.new(value, 0) unless node

    if value < node.value
      node.left = _put(value, node.left)
    else
      node.right = _put(value, node.right)
    end

    node.height = calculate_height(node)

    balance(node)
  end

  def calculate_height(node)
    1 + [height(node.left), height(node.right)].max
  end

  def balance_factor(node)
    height(node.left) - height(node.right)
  end

  def balance(node)
    # unbalanced left
    if balance_factor(node) > 1
      node.left = rotate_left(node.left) if balance_factor(node.left) < 0
      node = rotate_right(node)
    # unbalanced right
    elsif balance_factor(node) < -1
      node.right = rotate_right(node.right) if balance_factor(node.right) > 0
      node = rotate_left(node)
    end

    node
  end

  def rotate_right(node)
    y           = node.left
    node.left   = y.right
    y.right     = node
    node.height = calculate_height(node)
    y.height    = calculate_height(y)
    @root = y if node == root
    y
  end

  def rotate_left(node)
    y = node.right
    node.right = y.left
    y.left = node
    node.height = calculate_height(node)
    y.height    = calculate_height(y)
    @root = y if node == root
    y
  end

  def traverse_inorder(node = root, result = [])
    traverse_inorder(node.left, result) if node.left
    result << node.value
    traverse_inorder(node.right, result) if node.right

    result
  end

  def min(node = root)
    return node unless node.left

    min(node.left)
  end

  def delete_min(node = root)
    raise 'operation not permitted' unless @root

    _delete_min(node)
  end

  def _delete_min(node)
    return node.right unless node.left
    node.left   = _delete_min(node.left)
    node.height = calculate_height(node)

    balance(node)
  end

  def delete(value, node = root)
    if node && value < node.value
      node.left = delete(value, node.left)
    elsif node && value > node.value
      node.right = delete(value, node.right)
    else
      # no left child
      if node.left.nil?
        return node.right
      # no right child
      elsif node.right.nil?
        return node.left
      # children
      else
        # exchanges referecences with it's successor and delete the node
        y = node
        node = min(y.right) # successor
        node.right = delete_min(y.right)
        node.left  = y.left
      end
    end
  end

  def find(value, node = root)
    if node && value < node.value
      find(value, node.left)
    elsif node && value > node.value
      find(value, node.right)
    else
      node
    end
  end
end

avl_spec.rb

require 'rspec'
require './avl'

describe AvlTree do
  context '.initialize' do
    context 'empty tree' do
      subject { described_class.new }

      it 'has null root' do
        expect(subject.root).to be_nil
      end

      it 'height will be -1' do
        expect(subject.height).to eq(-1)
      end
    end
  end

  context '#put' do
    context 'empty tree' do
      subject { described_class.new }

      before { subject.put(10) }

      it 'creates a root node' do
        expect(subject.root).to_not be_nil
        expect(subject.root.value).to eq(10)
      end

      it 'height will be 0' do
        expect(subject.root.height).to eq(0)
      end

      it 'balance factor will be 0' do
        expect(subject.balance_factor(subject.root)).to eq(0)
      end
    end

    context 'no empty tree' do
      subject do
        tree = described_class.new
        tree.put(10)
        tree
      end

      context 'value smaller than current node' do
        before { subject.put(5) }

        it 'will be inserted on left side of current node' do
          expect(subject.root.left).to_not be_nil
        end

        it 'height will be 0' do
          expect(subject.root.left.height).to eq(0)
        end

        it 'parent node will have height 1' do
          expect(subject.root.height).to eq(1)
        end

        it 'balance factor will be 0' do
          expect(subject.balance_factor(subject.root.left)).to eq(0)
        end

        context 'unbalanced to left' do
          before { subject.put(2) }

          it 'will be rotated to right' do
            expect(subject.root.value).to eq(5)
            expect(subject.root.left.value).to eq(2)
            expect(subject.root.right.value).to eq(10)
          end
        end
      end

      context 'value bigger than current node' do
        before { subject.put(20) }

        it 'will be inserted on right side of current node' do
          expect(subject.root.right).to_not be_nil
        end

        it 'height will be 0' do
          expect(subject.root.right.height).to eq(0)
        end

        it 'parent node will have height 1' do
          expect(subject.root.height).to eq(1)
        end

        it 'balance factor will be 0' do
          expect(subject.balance_factor(subject.root.right)).to eq(0)
        end

        context 'unbalanced to right' do
          before { subject.put(30) }

          it 'will be rotated to left' do
            expect(subject.root.value).to eq(20)
            expect(subject.root.left.value).to eq(10)
            expect(subject.root.right.value).to eq(30)
          end

          it 'has 10, 20, 30' do
            expect(subject.traverse_inorder(subject.root)).to eq([10, 20, 30])
          end
        end
      end
    end
  end

  context '#find' do
    context 'empty' do
      subject { described_class.new }

      it 'returns nil' do
        expect(subject.find(10)).to be_nil
      end
    end

    context 'root' do
      subject { described_class.new }

      before { subject.put(10) }

      it 'returns root' do
        expect(subject.find(10)).to eq(subject.root)
      end
    end

    context 'with more than one item' do
      subject { described_class.new }

      before do
        subject.put(10)
        subject.put(20)
      end

      it 'returns 20 node' do
        expected = subject.find(20)
        expect(expected.value).to eq(20)
      end
    end
  end

  describe '#min' do
    context 'root' do
      subject { described_class.new }

      before { subject.put(10) }

      it 'returns root' do
        expect(subject.min).to eq(subject.root)
      end
    end

    context 'more than one item' do
      subject { described_class.new }

      before do
        subject.put(10)
        subject.put(20)
        subject.put(5)
        subject.put(40)
        subject.put(2)
      end

      it 'returns min node' do
        expected = subject.min
        expect(expected.value).to eq(2)
      end
    end
  end

  describe '#delete_min' do
    context 'empty tree' do
      subject { described_class.new }

      it 'raises error' do
        expect { subject.delete_min }.to raise_error 'operation not permitted'
      end
    end

    context 'not empty tree' do
      subject { described_class.new }

      before do
        subject.put(10)
        subject.put(20)
        subject.put(5)
        subject.put(40)
        subject.put(2)
      end

      it 'deletes the node 2' do
        subject.delete_min

        expect(subject.traverse_inorder).to eq([5, 10, 20, 40])
      end
    end
  end

  describe '#delete' do
    subject { described_class.new }

    before do
      subject.put(10)
      subject.put(20)
      subject.put(5)
      subject.put(40)
      subject.put(2)
    end

    it 'deletes the node' do
      subject.delete(5)

      expect(subject.traverse_inorder).to eq([2, 10, 20, 40])
    end

    it 'keeps the tree balanced' do
      subject.delete(5)

      expect(subject.balance_factor(subject.root)).to eq(-1)
    end
  end
end