aherrman · January 24, 2010 02:24 · Jan 24, 2010 · Jan 24, 2010 · Jan 24, 2010
diff --git a/sorted_queue.rb b/sorted_queue.rb
@@ -23,23 +23,31 @@ def to_a
     }
   end
 
+  def to_s
+    to_a.to_s
+  end
+
+  def inspect
+    to_a.inspect
+  end
+
   # Adds an item to the queue
-  def add(item)
+  def add!(item)
     @queue[item] = item
   end
 
   # Adds an item to the queue
   def <<(item)
-    add(item)
+    add!(item)
   end
 
   # Removes an item from the queue
-  def remove(item)
+  def remove!(item)
     @queue.delete(item)
   end
 
   # Removes the first element from the queue
-  def shift
+  def shift!
     shifted = @queue.shift
 
     return nil if shifted.nil?

diff --git a/gistfile1.rb → sorted_queue.rb b/gistfile1.rb → sorted_queue.rb
diff --git a/gistfile1.rb b/gistfile1.rb
@@ -0,0 +1,60 @@
+require 'rbtree'
+
+# A queue that keeps its elements sorted based on the <=> operator.
+# Requires the rbtree gem
+class SortedQueue
+  include Enumerable
+
+  def initialize(ary=[])
+    @queue = RBTree.new
+
+    ary.each { |item| 
+      @queue[item] = item
+    }
+  end
+
+  def size
+    @queue.size
+  end
+
+  def to_a
+    @queue.to_a.inject([]) { |a, keyval|
+      a << keyval[0]
+    }
+  end
+
+  # Adds an item to the queue
+  def add(item)
+    @queue[item] = item
+  end
+
+  # Adds an item to the queue
+  def <<(item)
+    add(item)
+  end
+
+  # Removes an item from the queue
+  def remove(item)
+    @queue.delete(item)
+  end
+
+  # Removes the first element from the queue
+  def shift
+    shifted = @queue.shift
+
+    return nil if shifted.nil?
+    shifted[0]
+  end
+
+  # Checks if the item is contained in the queue
+  def include?(item)
+    @queue.include?(item)
+  end
+
+  def each
+    @queue.each_key { |key|
+      yield key
+    }
+    self
+  end
+end