multivariate_hypergeometric_distribution.rb

module Math
  
  # https://rosettacode.org/wiki/Evaluate_binomial_coefficients#Ruby
  # binomial coefficient
	def self.bc(n, k)
		pTop = (n-k+1 .. n).inject(1, &:*) 
		pBottom = (2 .. k).inject(1, &:*)
		pTop / pBottom
	end
end

class MHD
	include Math

	attr_reader :population_size, :marble_distribution

	def initialize(marble_distribution: nil)
		@marble_distribution = marble_distribution
		@population_size = marble_distribution.values.inject(:+)
	end

	def call(amounts_desired: amounts_desired, draws: draws)
		raise "specify the amount desired for each type" unless amounts_desired.size == marble_distribution.size

		number_of_type_to_number_desired = marble_distribution.values.zip(amounts_desired)

		binomial_coefficients = number_of_type_to_number_desired.map do |number_of_type, number_desired|
			Math.bc(number_of_type, number_desired)
		end

		binomial_coefficients.inject(:*).to_f / Math.bc(population_size, draws).to_f
	end
end

p "does binomial work?"
p Math.bc(40, 4) == 91390

# a jar of 30 marbles consists of 5 red, 10 blue, and 15 yellow marbles
# a random sample of 6 marbles is chosen
marble_distribution = {
	red_marbles: 5,
	blue_marbles: 10,
	yellow_marbles: 15
}

mhd = MHD.new(marble_distribution: marble_distribution)

p "the probability that the sample contains exactly 2 red, 2 blue, and 2 yellow marbles is:"
p probability = mhd.call(amounts_desired: [2, 2, 2], draws: 6)

p "MHD works?"
p probability.round(4) == 0.0796

# https://en.wikipedia.org/wiki/Hypergeometric_distribution
# http://www.math.uah.edu/stat/dist/Discrete.html