Fetching contributors…
Cannot retrieve contributors at this time
334 lines (295 sloc) 6.55 KB
 # Algorithm based on "Introduction to Algorithms" by Cormen and others class RedBlackTree class Node attr_accessor :color attr_accessor :key attr_accessor :left attr_accessor :right attr_accessor :parent RED = :red BLACK = :black COLORS = [RED, BLACK].freeze def initialize(key, color = RED) raise ArgumentError, "Bad value for color parameter" unless COLORS.include?(color) @color = color @key = key @left = @right = @parent = NilNode.instance end def black? return color == BLACK end def red? return color == RED end end class NilNode < Node class << self private :new @instance = nil # it's not thread safe def instance if @instance.nil? @instance = new def instance return @instance end end return @instance end end def initialize self.color = BLACK self.key = 0 self.left = nil self.right = nil self.parent = nil end def nil? return true end end include Enumerable attr_accessor :root attr_accessor :size def initialize self.root = NilNode.instance self.size = 0 end def add(key) insert(Node.new(key)) end def insert(x) insert_helper(x) x.color = Node::RED while x != root && x.parent.color == Node::RED if x.parent == x.parent.parent.left y = x.parent.parent.right if !y.nil? && y.color == Node::RED x.parent.color = Node::BLACK y.color = Node::BLACK x.parent.parent.color = Node::RED x = x.parent.parent else if x == x.parent.right x = x.parent left_rotate(x) end x.parent.color = Node::BLACK x.parent.parent.color = Node::RED right_rotate(x.parent.parent) end else y = x.parent.parent.left if !y.nil? && y.color == Node::RED x.parent.color = Node::BLACK y.color = Node::BLACK x.parent.parent.color = Node::RED x = x.parent.parent else if x == x.parent.left x = x.parent right_rotate(x) end x.parent.color = Node::BLACK x.parent.parent.color = Node::RED left_rotate(x.parent.parent) end end end root.color = Node::BLACK end alias << insert def delete(z) y = (z.left.nil? || z.right.nil?) ? z : successor(z) x = y.left.nil? ? y.right : y.left x.parent = y.parent if y.parent.nil? self.root = x else if y == y.parent.left y.parent.left = x else y.parent.right = x end end z.key = y.key if y != z if y.color == Node::BLACK delete_fixup(x) end self.size -= 1 return y end def minimum(x = root) while !x.left.nil? x = x.left end return x end def maximum(x = root) while !x.right.nil? x = x.right end return x end def successor(x) if !x.right.nil? return minimum(x.right) end y = x.parent while !y.nil? && x == y.right x = y y = y.parent end return y end def predecessor(x) if !x.left.nil? return maximum(x.left) end y = x.parent while !y.nil? && x == y.left x = y y = y.parent end return y end def inorder_walk(x = root) x = self.minimum while !x.nil? yield x.key x = successor(x) end end alias each inorder_walk def reverse_inorder_walk(x = root) x = self.maximum while !x.nil? yield x.key x = predecessor(x) end end alias reverse_each reverse_inorder_walk def search(key, x = root) while !x.nil? && x.key != key key < x.key ? x = x.left : x = x.right end return x end def empty? return self.root.nil? end def black_height(x = root) height = 0 while !x.nil? x = x.left height +=1 if x.nil? || x.black? end return height end private def left_rotate(x) raise "x.right is nil!" if x.right.nil? y = x.right x.right = y.left y.left.parent = x if !y.left.nil? y.parent = x.parent if x.parent.nil? self.root = y else if x == x.parent.left x.parent.left = y else x.parent.right = y end end y.left = x x.parent = y end def right_rotate(x) raise "x.left is nil!" if x.left.nil? y = x.left x.left = y.right y.right.parent = x if !y.right.nil? y.parent = x.parent if x.parent.nil? self.root = y else if x == x.parent.left x.parent.left = y else x.parent.right = y end end y.right = x x.parent = y end def insert_helper(z) y = NilNode.instance x = root while !x.nil? y = x z.key < x.key ? x = x.left : x = x.right end z.parent = y if y.nil? self.root = z else z.key < y.key ? y.left = z : y.right = z end self.size += 1 end def delete_fixup(x) while x != root && x.color == Node::BLACK if x == x.parent.left w = x.parent.right if w.color == Node::RED w.color = Node::BLACK x.parent.color = Node::RED left_rotate(x.parent) w = x.parent.right end if w.left.color == Node::BLACK && w.right.color == Node::BLACK w.color = Node::RED x = x.parent else if w.right.color == Node::BLACK w.left.color = Node::BLACK w.color = Node::RED right_rotate(w) w = x.parent.right end w.color = x.parent.color x.parent.color = Node::BLACK w.right.color = Node::BLACK left_rotate(x.parent) x = root end else w = x.parent.left if w.color == Node::RED w.color = Node::BLACK x.parent.color = Node::RED right_rotate(x.parent) w = x.parent.left end if w.right.color == Node::BLACK && w.left.color == Node::BLACK w.color = Node::RED x = x.parent else if w.left.color == Node::BLACK w.right.color = Node::BLACK w.color = Node::RED left_rotate(w) w = x.parent.left end w.color = x.parent.color x.parent.color = Node::BLACK w.left.color = Node::BLACK right_rotate(x.parent) x = root end end end x.color = Node::BLACK end end