roulette-wheel-selection.rb in roulette-wheel-selection-1.0.2

- old
+ new

@@ -17,16 +17,44 @@
       hash = array.map{|v| [v, v[key]] }.to_h
       return sample_from_hash(hash)
     end
 
     def sample_from_hash(hash)
-      total_rate = hash.values.inject(&:+) || 0
-      return if total_rate == 0
-      random_seed = rand(total_rate)
-      hash.each do |obj, rate|
-        return obj if random_seed < rate
-        random_seed -= rate
-      end
+      RouletteWheelSelection.new(hash).sample
+    end
+  end
+
+  def initialize(hash)
+    @hash = hash
+    @total_rate = hash.values.inject(&:+) || 0
+  end
+
+  def sample(num = 1)
+    return if @total_rate == 0
+    return if num < 1
+    return sample_n_objects(num) if num > 1
+    return sample_an_object(@total_rate, @hash)
+  end
+
+  private
+
+  def sample_n_objects(num)
+    hash = @hash.clone
+    total_rate = @total_rate
+    num = total_rate if num > total_rate
+    return Array.new(num) do
+      obj = sample_an_object(total_rate, hash)
+      hash[obj] -= 1
+      total_rate -= 1
+      next obj
+    end
+  end
+
+  def sample_an_object(total_rate, hash)
+    random_seed = rand(total_rate)
+    hash.each do |obj, rate|
+      return obj if random_seed < rate
+      random_seed -= rate
     end
   end
 end
 