rbtree.rb

module RBTree
  module M
    def insert val
      RBTree.insert self, val
    end
    alias :<< :insert
    def include?(...) = key?(...)
  end

  class Leaf
    include M

    def color     = :black
    def key?(_)   = false

    def deconstruct
      [:black, nil, nil, nil]
    end
  end

  class Node
    include M

    attr_reader :color, :key, :left, :right

    def initialize color, key, left, right
      @color = color
      @key   = key
      @left  = left
      @right = right
    end

    def deconstruct
      [color, key, left, right]
    end

    def key? k
      k == key || (k < key ? left.key?(k) : right.key?(k))
    end
  end

  LEAF = Leaf.new

  def self.new; LEAF; end

  def self.insert tree, key
    root = insert_aux(tree, key)
    case root
    in [:red, key, left, right]
      Node.new(:black, key, left, right)
    else
      root
    end
  end

  private_class_method def self.balance color, key, left, right
    # x, y, z are keys
    # a, b, c, d are nodes
    case [color, key, left, right]
    #   HEAD        LEFT      LEFT-LEFT
    in [:black, z, [:red, y, [:red, x, a, b], c], d]
    #   HEAD        LEFT         LEFT-RIGHT
    in [:black, z, [:red, x, a, [:red, y, b, c]], d]
    #   HEAD           RIGHT        RIGHT-RIGHT
    in [:black, x, a, [:red, y, b, [:red, z, c, d]]]
    #   HEAD           RIGHT     RIGHT-LEFT
    in [:black, x, a, [:red, z, [:red, y, b, c], d]]

    else
      return Node.new(color, key, left, right)
    end

    Node.new(:red, y, Node.new(:black, x, a, b), Node.new(:black, z, c, d))
  end

  private_class_method def self.insert_aux tree, key
    if tree == LEAF
      Node.new :red, key, LEAF, LEAF
    else
      # If the new key is smaller, we need to put it on the left
      if key < tree.key
        new_left = insert_aux(tree.left, key)

        balance(tree.color, tree.key, new_left, tree.right)

      # If the new key is larger, we need to put it on the right
      elsif key > tree.key
        new_right = insert_aux(tree.right, key)

        balance(tree.color, tree.key, tree.left, new_right)
      else
        tree
      end
    end
  end
end

if __FILE__ == $0
  tree = RBTree.new # empty
  tree1 =  tree << 5
  tree2 = tree1 << 3
  tree3 = tree2 << 2

  p tree1.include?(2) # => false
  p tree3.include?(2) # => true
end