Merge pull request #22 from ngng628/priority-queue

Improve PriorityQueue

spec/priority_queue_spec.cr

@@ -41,6 +41,22 @@ describe "PriorityQueue" do
     q.pop.should eq nil
   end
 
+  describe ".pop!" do
+    it "pops values in priority order. Raises Enumerable::NotFoundError if queue is empty." do
+      q = PriorityQueue(Int32).max
+      q << 5 << 6 << 1 << 3 << 2 << 8 << 7 << 4
+      q.pop!.should eq 8
+      q.pop!.should eq 7
+      q.pop!.should eq 6
+      q.pop!.should eq 5
+      q.pop!.should eq 4
+      q.pop!.should eq 3
+      q.pop!.should eq 2
+      q.pop!.should eq 1
+      expect_raises(Enumerable::EmptyError) { q.pop! }
+    end
+  end
+
   it "initialized by elements in *enumerable*" do
     q = PriorityQueue.max([5, 6, 1, 3, 2, 8, 7, 4])
     q.pop.should eq 8

src/priority_queue.cr

@@ -114,7 +114,8 @@ module AtCoder
     end
 
     # Pushes value into the queue.
-    def push(v : T)
+    # This method returns self, so several calls can be chained.
+    def push(v : T) : self
       @heap << v
       index = @heap.size - 1
       while index != 0
@@ -125,10 +126,11 @@ module AtCoder
         @heap[parent], @heap[index] = @heap[index], @heap[parent]
         index = parent
       end
+      self
     end
 
     # Alias of `push`
-    def <<(v : T)
+    def <<(v : T) : self
       push(v)
     end
 
@@ -158,6 +160,12 @@ module AtCoder
       ret
     end
 
+    # Pops value from the queue.
+    # Raises `Enumerable::EmptyError` if queue is of 0 size.
+    def pop!
+      pop || raise Enumerable::EmptyError.new
+    end
+
     # Yields each item in the queue in comparator's order.
     def each(&)
       @heap.sort { |a, b| @compare_proc.call(a, b) ? 1 : -1 }.each do |e|