khiav223577/roulette-wheel-selection

# 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)
      when Hash  ; sample_from_hash(*args)
      else       ; fail "Unsupported type: #{object.class}"
      end
    end

    private

    def sample_from_array(array, key)
      hash = array.map{|v| [v, v[key]] }.to_h
      return sample_from_hash(hash)
    end

    def sample_from_hash(hash)
      RouletteWheelSelection.new(hash).sample
    end
  end

  def initialize(hash)
    @hash = hash
    @total_rate = hash.values.inject(0, :+)
  end

  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)]
  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