Skip to content

Instantly share code, notes, and snippets.

@chtzvt
Forked from tenderlove/rb_tree.rb
Created January 31, 2023 19:02
Show Gist options
  • Save chtzvt/73461bff63e51c4666b57cd3d0626d93 to your computer and use it in GitHub Desktop.
Save chtzvt/73461bff63e51c4666b57cd3d0626d93 to your computer and use it in GitHub Desktop.
# 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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment