Simulating a Loaded Dice in a Constant Time

alias_method.rb

#
# Implemented from http://web.eecs.utk.edu/~vose/Publications/random.pdf
# and http://scriptogr.am/jj/post/simulating-a-loaded-dice-in-a-constant-time
#

class AliasMethod

   def initialize(probs)
      @probability = probs
      @alias_numbers = [ -1 ] * @probability.size
      avg = 1.0 / @probability.size

      large, small = (0..(@probability.size - 1)).zip(@probability).partition { |i, p| p >= avg }

      large.collect! { |i,p| i }
      small.collect! { |i,p| i }

      while !small.empty? and !large.empty? do
         less = small.slice! -1
         more = large.slice! -1

         @alias_numbers[less] = more

         @probability[more] = (@probability[more] + @probability[less]) - avg

         if (@probability[more] >= avg)
            large = large + [more]
         else
            small = small + [more]
         end
      end

      small.each { |p| @probability[p] = avg }
      large.each { |p| @probability[p] = avg }

      @probability.map! { |p| p * @probability.size }
   end

   def next
      col = Random.rand(@probability.size)
      coinToss = Random.rand(0.0...1.0) < @probability[col]
      coinToss ? col : @alias_numbers[col]
   end

end