rb_tree.rb

# Okasaki style Functional Red Black Tree
# https://www.cs.tufts.edu/comp/150FP/archive/chris-okasaki/redblack99.pdf
#
# leaves and root are black
# BST
# No red node has a red child
# Every path from the root to a leaf has the same number of black nodes

module RBTree
  class Leaf
    include Enumerable

    def color     = :black
    def leaf?     = true
    def key?(_)   = false
    def each; end

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

    def insert val
      RBTree.insert self, val
    end
  end

  class Node
    include Enumerable

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

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

    def leaf? = false

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

    def key? k
      return true if k == key
      k < key ? left.key?(k) : right.key?(k)
    end

    def insert val
      RBTree.insert self, val
    end

    def each &blk
      yield self
      left.each(&blk)
      right.each(&blk)
    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]
    in [:black, z, [:red, y, [:red, x, a, b], c], d]
    in [:black, z, [:red, x, a, [:red, y, b, c]], d]
    in [:black, x, a, [:red, z, [:red, y, b, c], d]]
    in [:black, x, a, [:red, y, b, [:red, z, 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, x
    if tree.leaf?
      Node.new :red, x, LEAF, LEAF
    else
      if x < tree.key
        balance(tree.color, tree.key, insert_aux(tree.left, x), tree.right)
      elsif x > tree.key
        balance(tree.color, tree.key, tree.left, insert_aux(tree.right, x))
      else
        tree
      end
    end
  end

  def self.insert tree, x
    root = insert_aux(tree, x)
    case root
    in [:red, l, v, r]
      Node.new(:black, l, v, r)
    else
      root
    end
  end

  LEAF = Leaf.new

  def self.new; LEAF; end
end

if $0 == __FILE__
  tree = RBTree.new
  tree = tree.insert 5
  tree = tree.insert 4
  tree = tree.insert 2
  tree = tree.insert 7
  tree = tree.insert 8
  tree = tree.insert 9
  tree = tree.insert 1
  tree = tree.insert 3
  5000.times { tree = tree.insert rand(123456) }
  p tree.key?(8)  # true
  p tree.key?(10) # maybe false
end