lib/roulette-wheel-selection.rb in roulette-wheel-selection-1.0.2 vs lib/roulette-wheel-selection.rb in roulette-wheel-selection-1.1.0

@@ -1,8 +1,12 @@
+# frozen_string_literal: true
+
 require "roulette-wheel-selection/version"
 
 class RouletteWheelSelection
+  NOT_SET = Object.new
+
   class << self
     def sample(*args)
       object = args.first
       case object
       when Array ; sample_from_array(*args)
@@ -26,14 +30,15 @@
   def initialize(hash)
     @hash = hash
     @total_rate = hash.values.inject(&:+) || 0
   end
 
-  def sample(num = 1)
+  def sample(num = NOT_SET)
     return if @total_rate == 0
+    return sample_an_object(@total_rate, @hash) if num == NOT_SET
     return if num < 1
     return sample_n_objects(num) if num > 1
-    return sample_an_object(@total_rate, @hash)
+    return [sample_an_object(@total_rate, @hash)]
   end
 
   private
 
   def sample_n_objects(num)