Created
March 29, 2024 00:13
-
-
Save tenderlove/44fd4206b9cbc42bcafc90adb3f5820d to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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 |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment