avl_tree 1.0.0 → 1.1.0

data/bench/bench.rb

@@ -1,8 +1,10 @@
 require 'benchmark'
 require 'radix_tree' # gem install radix_tree
 require 'avl_tree'
+require 'red_black_tree'
+require 'openssl'
 
-random = Random.new(0)
+#random = Random.new(0)
 
 TIMES = 100000
 key_size = 10
@@ -40,11 +42,12 @@ end
 
 keys = []
 TIMES.times do
-  keys << random.bytes(key_size)
+  keys << OpenSSL::Random.random_bytes(key_size)
 end
 
 Benchmark.bmbm do |bm|
   run(bm, Hash.new, keys)
   run(bm, RadixTree.new, keys)
   run(bm, AVLTree.new, keys)
+  run(bm, RedBlackTree.new, keys)
 end

data/bench/bench_element_size.rb

@@ -1,8 +1,7 @@
 require 'benchmark'
 require 'radix_tree' # gem install radix_tree
 require 'avl_tree'
-
-random = Random.new(0)
+require 'openssl'
 
 times = 100000
 key_size = 10
@@ -38,15 +37,17 @@ def run(bm, h, keys)
   end
 end
 
-[10000, 20000, 50000, 100000, 200000, 500000].each do |elements|
-  keys = []
-  elements.times do
-    keys << random.bytes(key_size)
-  end
+keys = []
+1000000.times do
+  keys << OpenSSL::Random.random_bytes(key_size)
+end
+
+1.upto(100) do |idx|
+  elements = idx * 10000
 
   Benchmark.bm(30) do |bm|
-    run(bm, Hash.new, keys)
-    run(bm, RadixTree.new, keys)
-    run(bm, AVLTree.new, keys)
+    #run(bm, Hash.new, keys[0, elements])
+    #run(bm, RadixTree.new, keys)
+    run(bm, AVLTree.new, keys[0, elements])
   end
 end

data/bench/profile.rb

@@ -1,13 +1,13 @@
-require File.expand_path('../lib/avl_tree', File.dirname(__FILE__))
+require File.expand_path('../lib/red_black_tree', File.dirname(__FILE__))
 
 random = Random.new(0)
 
-TIMES = 200000
+TIMES = 50000
 key_size = 10
 
-h = AVLTree.new
+h = RedBlackTree.new
 TIMES.times do
   h[random.bytes(key_size)] = 1
-  h[random.bytes(key_size)]
-  h.delete(random.bytes(key_size))
+  #h[random.bytes(key_size)]
+  #h.delete(random.bytes(key_size))
 end

data/lib/avl_tree.rb

@@ -26,7 +26,7 @@ class AVLTree
       end
 
       # returns new_root
-      def store(key, value)
+      def insert(key, value)
         Node.new(key, value)
       end
 
@@ -86,13 +86,13 @@ class AVLTree
       yield [@key, @value]
       @right.each(&block)
     end
-    
+
     def each_key
       each do |k, v|
         yield k
       end
     end
-    
+
     def each_value
       each do |k, v|
         yield v
@@ -108,14 +108,14 @@ class AVLTree
     end
 
     # returns new_root
-    def store(key, value)
+    def insert(key, value)
       case key <=> @key
       when -1
-        @left = @left.store(key, value)
+        @left = @left.insert(key, value)
       when 0
         @value = value
       when 1
-        @right = @right.store(key, value)
+        @right = @right.insert(key, value)
       end
       rotate
     end
@@ -213,17 +213,15 @@ class AVLTree
     def rotate
       case @left.height - @right.height
       when +2
-        if @left.left.height >= @left.right.height
-          root = rotate_LL
-        else
-          root = rotate_LR
+        if @left.left.height < @left.right.height
+          @left = @left.rotate_left
         end
+        root = rotate_right
       when -2
-        if @right.left.height <= @right.right.height
-          root = rotate_RR
-        else
-          root = rotate_RL
+        if @right.left.height > @right.right.height
+          @right = @right.rotate_right
         end
+        root = rotate_left
       else
         root = self
       end
@@ -231,21 +229,6 @@ class AVLTree
       root
     end
 
-  private
-
-    def delete_self
-      if @left.empty? and @right.empty?
-        deleted = EMPTY
-      elsif @right.height < @left.height
-        deleted, new_left = @left.delete_max
-        deleted.left, deleted.right = new_left, @right
-      else
-        deleted, new_right = @right.delete_min
-        deleted.left, deleted.right = @left, new_right
-      end
-      deleted
-    end
-
     # Right single rotation
     # (B a (D c E)) where D-a > 1 && E > c --> (D (B a c) E)
     #
@@ -255,7 +238,7 @@ class AVLTree
     #    / \        / \
     #   c   E      a   c
     #
-    def rotate_RR
+    def rotate_left
       root = @right
       @right = root.left
       root.left = self
@@ -272,7 +255,7 @@ class AVLTree
     #  / \            / \
     # A   c          c   e
     #
-    def rotate_LL
+    def rotate_right
       root = @left
       @left = root.right
       root.right = self
@@ -280,50 +263,19 @@ class AVLTree
       root
     end
 
-    # Right double rotation
-    # (B a (F (D c e) g)) where F-a > 1 && D > g --> (D (B a c) (F e g))
-    #
-    #   B               D
-    #  / \            /   \
-    # a   F    ->    B     F
-    #    / \        / \   / \
-    #   D   g      a   c e   g
-    #  / \
-    # c   e
-    #
-    def rotate_RL
-      other = @right
-      root = other.left
-      @right = root.left
-      other.left = root.right
-      root.left = self
-      root.right = other
-      root.left.update_height
-      root.right.update_height
-      root
-    end
+  private
 
-    # Left double rotation
-    # (F (B a (D c e)) g) where B-g > 1 && D > a --> (d (B a c) (F e g))
-    #
-    #     F             D
-    #    / \          /   \
-    #   B   g  ->    B     F
-    #  / \          / \   / \
-    # a   D        a   c e   g
-    #    / \
-    #   c   e
-    #
-    def rotate_LR
-      other = @left
-      root = other.right
-      @left = root.right
-      other.right = root.left
-      root.right = self
-      root.left = other
-      root.left.update_height
-      root.right.update_height
-      root
+    def delete_self
+      if @left.empty? and @right.empty?
+        deleted = EMPTY
+      elsif @right.height < @left.height
+        deleted, new_left = @left.delete_max
+        deleted.left, deleted.right = new_left, @right
+      else
+        deleted, new_right = @right.delete_min
+        deleted.left, deleted.right = @left, new_right
+      end
+      deleted
     end
 
     def collect
@@ -339,7 +291,7 @@ class AVLTree
 
   attr_accessor :default
   attr_reader :default_proc
-  
+
   def initialize(default = DEFAULT, &block)
     if block && default != DEFAULT
       raise ArgumentError, 'wrong number of arguments'
@@ -403,9 +355,9 @@ class AVLTree
   end
 
   def []=(key, value)
-    @root = @root.store(key.to_s, value)
+    @root = @root.insert(key.to_s, value)
   end
-  alias store []=
+  alias insert []=
 
   def key?(key)
     @root.retrieve(key.to_s) != Node::UNDEFINED

data/lib/red_black_tree.rb

@@ -0,0 +1,549 @@
+class RedBlackTree
+  include Enumerable
+
+  class Node
+    UNDEFINED = Object.new
+
+    attr_reader :key, :value, :color
+    attr_reader :left, :right
+
+    def initialize(key, value)
+      @key, @value = key, value
+      @left = @right = EMPTY
+      # new node is added as RED
+      @color = :RED
+    end
+
+    def set_root
+      @color = :BLACK
+    end
+
+    def red?
+      @color == :RED
+    end
+
+    def black?
+      @color == :BLACK
+    end
+
+    def empty?
+      false
+    end
+
+    def size
+      @left.size + 1 + @right.size
+    end
+
+    # inorder
+    def each(&block)
+      @left.each(&block)
+      yield [@key, @value]
+      @right.each(&block)
+    end
+
+    def each_key
+      each do |k, v|
+        yield k
+      end
+    end
+
+    def each_value
+      each do |k, v|
+        yield v
+      end
+    end
+
+    def keys
+      collect { |k, v| k }
+    end
+
+    def values
+      collect { |k, v| v }
+    end
+
+    # returns new_root
+    def insert(key, value)
+      ret = self
+      case key <=> @key
+      when -1
+        @left = @left.insert(key, value)
+        if black? and @left.red? and !@left.children_both_black?
+          ret = rebalance_for_left_insert
+        end
+      when 0
+        @value = value
+      when 1
+        @right = @right.insert(key, value)
+        if black? and @right.red? and !@right.children_both_black?
+          ret = rebalance_for_right_insert
+        end
+      end
+      ret
+    end
+
+    # returns value
+    def retrieve(key)
+      case key <=> @key
+      when -1
+        @left.retrieve(key)
+      when 0
+        @value
+      when 1
+        @right.retrieve(key)
+      end
+    end
+
+    # returns [deleted_node, new_root, is_rebalance_needed]
+    def delete(key)
+      ret = self
+      case key <=> @key
+      when -1
+        deleted, @left, rebalance = @left.delete(key)
+        if rebalance
+          ret, rebalance = rebalance_for_left_delete
+        end
+      when 0
+        deleted = self
+        ret, rebalance = delete_self
+      when 1
+        deleted, @right, rebalance = @right.delete(key)
+        if rebalance
+          ret, rebalance = rebalance_for_right_delete
+        end
+      end
+      [deleted, ret, rebalance]
+    end
+
+    def dump_tree(io, indent = '')
+      @right.dump_tree(io, indent + '  ')
+      io << indent << sprintf("#<%s:0x%010x %s %s> => %s", self.class.name, __id__, @color, @key.inspect, @value.inspect) << $/
+      @left.dump_tree(io, indent + '  ')
+    end
+
+    def dump_sexp
+      left = @left.dump_sexp
+      right = @right.dump_sexp
+      if left or right
+        '(' + [@key, left || '-', right].compact.join(' ') + ')'
+      else
+        @key
+      end
+    end
+
+    # for debugging
+    def check_height
+      lh = @left.empty? ? 0 : @left.check_height
+      rh = @right.empty? ? 0 : @right.check_height
+      if red?
+        if @left.red? or @right.red?
+          puts dump_tree(STDERR)
+          raise 'red/red assertion failed'
+        end
+      else
+        if lh != rh
+          puts dump_tree(STDERR)
+          raise "black height unbalanced: #{lh} #{rh}"
+        end
+      end
+      (lh > rh ? lh : rh) + (black? ? 1 : 0)
+    end
+
+  protected
+
+    def children_both_black?
+      @right.black? and @left.black?
+    end
+
+    def color=(color)
+      @color = color
+    end
+
+    def left=(left)
+      @left = left
+    end
+
+    def right=(right)
+      @right = right
+    end
+
+    def color_flip(other)
+      @color, other.color = other.color, @color
+    end
+
+    def node_flip(other)
+      @left, other.left = other.left, @left
+      @right, other.right = other.right, @right
+      color_flip(other)
+    end
+
+    def delete_min
+      if @left.empty?
+        [self, *delete_self]
+      else
+        ret = self
+        deleted, @left, rebalance = @left.delete_min
+        if rebalance
+          ret, rebalance = rebalance_for_left_delete
+        end
+        [deleted, ret, rebalance]
+      end
+    end
+
+    # trying to rebalance when the left sub-tree is 1 level lower than the right
+    def rebalance_for_left_delete
+      ret = self
+      rebalance = false
+      if black?
+        if @right.black?
+          if @right.children_both_black?
+            # make whole sub-tree 1 level lower and ask rebalance
+            @right.color = :RED
+            rebalance = true
+          else
+            # move 1 black from the right to the left by single/double rotation
+            ret = balanced_rotate_left
+          end
+        else
+          # flip this sub-tree into another type of 3-children node
+          ret = rotate_left
+          # try to rebalance in sub-tree
+          ret.left, rebalance = ret.left.rebalance_for_left_delete
+          raise 'should not happen' if rebalance
+        end
+      else # red
+        if @right.children_both_black?
+          # make right sub-tree 1 level lower
+          color_flip(@right)
+        else
+          # move 1 black from the right to the left by single/double rotation
+          ret = balanced_rotate_left
+        end
+      end
+      [ret, rebalance]
+    end
+
+    # trying to rebalance when the right sub-tree is 1 level lower than the left
+    # See rebalance_for_left_delete.
+    def rebalance_for_right_delete
+      ret = self
+      rebalance = false
+      if black?
+        if @left.black?
+          if @left.children_both_black?
+            @left.color = :RED
+            rebalance = true
+          else
+            ret = balanced_rotate_right
+          end
+        else
+          ret = rotate_right
+          ret.right, rebalance = ret.right.rebalance_for_right_delete
+          raise 'should not happen' if rebalance
+        end
+      else # red
+        if @left.children_both_black?
+          color_flip(@left)
+        else
+          ret = balanced_rotate_right
+        end
+      end
+      [ret, rebalance]
+    end
+
+    # move 1 black from the right to the left by single/double rotation
+    def balanced_rotate_left
+      if @right.left.red? and @right.right.black?
+        @right = @right.rotate_right
+      end
+      ret = rotate_left
+      ret.right.color = ret.left.color = :BLACK
+      ret
+    end
+
+    # move 1 black from the left to the right by single/double rotation
+    def balanced_rotate_right
+      if @left.right.red? and @left.left.black?
+        @left = @left.rotate_left
+      end
+      ret = rotate_right
+      ret.right.color = ret.left.color = :BLACK
+      ret
+    end
+
+    # Right single rotation
+    # (b a (D c E)) where D and E are RED --> (d (B a c) E)
+    #
+    #   b              d
+    #  / \            / \
+    # a   D    ->    B   E
+    #    / \        / \
+    #   c   E      a   c
+    #
+    def rotate_left
+      root = @right
+      @right = root.left
+      root.left = self
+      root.color_flip(root.left)
+      root
+    end
+
+    # Left single rotation
+    # (d (B A c) e) where A and B are RED --> (b A (D c e))
+    #
+    #     d          b
+    #    / \        / \
+    #   B   e  ->  A   D
+    #  / \            / \
+    # A   c          c   e
+    #
+    def rotate_right
+      root = @left
+      @left = root.right
+      root.right = self
+      root.color_flip(root.right)
+      root
+    end
+
+  private
+
+    # trying to rebalance when the left sub-tree is 1 level higher than the right
+    def rebalance_for_left_insert
+      ret = self
+      if @right.red?
+        @color = :RED
+        @left.color = @right.color = :BLACK
+      else
+        if @left.right.red?
+          @left = @left.rotate_left
+        end
+        ret = rotate_right
+      end
+      ret
+    end
+
+    # trying to rebalance when the right sub-tree is 1 level higher than the left
+    def rebalance_for_right_insert
+      ret = self
+      if @left.red?
+        @color = :RED
+        @left.color = @right.color = :BLACK
+      else
+        if @right.left.red?
+          @right = @right.rotate_right
+        end
+        ret = rotate_left
+      end
+      ret
+    end
+
+    def delete_self
+      rebalance = false
+      if @left.empty? and @right.empty?
+        # just remove this node and ask rebalance to the parent
+        new_root = EMPTY
+        if black?
+          rebalance = true
+        end
+      elsif @left.empty? or @right.empty?
+        # pick the single children
+        new_root = @left.empty? ? @right : @left
+        if black?
+          # keep the color black
+          raise 'should not happen' unless new_root.red?
+          color_flip(new_root)
+        else
+          # just remove the red node
+        end
+      else
+        # pick the minimum node from the right sub-tree and replace self with it
+        new_root, @right, rebalance = @right.delete_min
+        new_root.node_flip(self)
+        if rebalance
+          new_root, rebalance = new_root.rebalance_for_right_delete
+        end
+      end
+      [new_root, rebalance]
+    end
+
+    def collect
+      pool = []
+      each do |key, value|
+        pool << yield(key, value)
+      end
+      pool
+    end
+
+    class EmptyNode
+      def red?
+        false
+      end
+
+      def black?
+        true
+      end
+
+      def empty?
+        true
+      end
+
+      def value
+        nil
+      end
+
+      def size
+        0
+      end
+
+      def each(&block)
+        # intentionally blank
+      end
+
+      # returns new_root
+      def insert(key, value)
+        Node.new(key, value)
+      end
+
+      # returns value
+      def retrieve(key)
+        UNDEFINED
+      end
+
+      # returns [deleted_node, new_root, is_rebalance_needed]
+      def delete(key)
+        [self, self, false]
+      end
+
+      def dump_tree(io, indent = '')
+        # intentionally blank
+      end
+
+      def dump_sexp
+        # intentionally blank
+      end
+    end
+    EMPTY = Node::EmptyNode.new
+  end
+
+  DEFAULT = Object.new
+
+  attr_accessor :default
+  attr_reader :default_proc
+
+  def initialize(default = DEFAULT, &block)
+    if block && default != DEFAULT
+      raise ArgumentError, 'wrong number of arguments'
+    end
+    @root = Node::EMPTY
+    @default = default
+    @default_proc = block
+  end
+
+  def empty?
+    @root == Node::EMPTY
+  end
+
+  def size
+    @root.size
+  end
+  alias length size
+
+  def each(&block)
+    if block_given?
+      @root.each(&block)
+      self
+    else
+      Enumerator.new(@root)
+    end
+  end
+  alias each_pair each
+
+  def each_key
+    if block_given?
+      @root.each do |k, v|
+        yield k
+      end
+      self
+    else
+      Enumerator.new(@root, :each_key)
+    end
+  end
+
+  def each_value
+    if block_given?
+      @root.each do |k, v|
+        yield v
+      end
+      self
+    else
+      Enumerator.new(@root, :each_value)
+    end
+  end
+
+  def keys
+    @root.keys
+  end
+
+  def values
+    @root.values
+  end
+
+  def clear
+    @root = Node::EMPTY
+  end
+
+  def []=(key, value)
+    @root = @root.insert(key.to_s, value)
+    @root.set_root
+    @root.check_height if $DEBUG
+  end
+  alias insert []=
+
+  def key?(key)
+    @root.retrieve(key.to_s) != Node::UNDEFINED
+  end
+  alias has_key? key?
+
+  def [](key)
+    value = @root.retrieve(key.to_s)
+    if value == Node::UNDEFINED
+      default_value
+    else
+      value
+    end
+  end
+
+  def delete(key)
+    deleted, @root, rebalance = @root.delete(key.to_s)
+    unless empty?
+      @root.set_root
+      @root.check_height if $DEBUG
+    end
+    deleted.value
+  end
+
+  def dump_tree(io = '')
+    @root.dump_tree(io)
+    io << $/
+    io
+  end
+
+  def dump_sexp
+    @root.dump_sexp || ''
+  end
+
+  def to_hash
+    inject({}) { |r, (k, v)| r[k] = v; r }
+  end
+
+private
+
+  def default_value
+    if @default != DEFAULT
+      @default
+    elsif @default_proc
+      @default_proc.call
+    else
+      nil
+    end
+  end
+end