Cacher is the code snippet organizer for pro developers

require 'rubygems'
require 'randomext'

class ThompsonSampling

  def initialize(sf_counts)
    @s = []
    @f = []
    sf_counts.each do |sf| 
      @s.push(sf[:s])
      @f.push(sf[:f])
    end
    @alpha = 1
    @beta = 1
    @num_of_arm = sf_counts.length
  end

  def distribute()
    distr = Array.new(@num_of_arm, 0)
    rand = Random.new()

    num_of_trial = 10000
    num_of_trial.times do
      params = []
      for idx in 1..@num_of_arm do
        param = {
          key: idx - 1,
          theta: rand.beta(@alpha + @s[idx - 1], @beta + @f[idx - 1])
        }
        params.push(param)
      end

      max_param = params.max_by {|param| param[:theta]}
      distr[max_param[:key]] += 1
    end

    p distr
    p distr.map! {|item| item * 100 / num_of_trial}
  end
end

sf_counts = [
  {s: 11, f:100},
  {s: 6, f:100},
  {s: 10, f:100}
]

ts = ThompsonSampling.new(sf_counts)
ts.distribute()