diff --git a/maths/binomial_distribution.py b/maths/binomial_distribution.py index e29c8597a6e5..71344f5e2afd 100644 --- a/maths/binomial_distribution.py +++ b/maths/binomial_distribution.py @@ -1,41 +1,43 @@ -"""For more information about the Binomial Distribution - -https://en.wikipedia.org/wiki/Binomial_distribution""" - -from math import factorial - - -def binomial_distribution(successes: int, trials: int, prob: float) -> float: - """ - Return probability of k successes out of n tries, with p probability for one - success - - The function uses the factorial function in order to calculate the binomial - coefficient - - >>> binomial_distribution(3, 5, 0.7) - 0.30870000000000003 - >>> binomial_distribution (2, 4, 0.5) - 0.375 - """ - if successes > trials: - raise ValueError("""successes must be lower or equal to trials""") - if trials < 0 or successes < 0: - raise ValueError("the function is defined for non-negative integers") - if not isinstance(successes, int) or not isinstance(trials, int): - raise ValueError("the function is defined for non-negative integers") - if not 0 < prob < 1: - raise ValueError("prob has to be in range of 1 - 0") - probability = (prob**successes) * ((1 - prob) ** (trials - successes)) - # Calculate the binomial coefficient: n! / k!(n-k)! - coefficient = float(factorial(trials)) - coefficient /= factorial(successes) * factorial(trials - successes) - return probability * coefficient - - -if __name__ == "__main__": - from doctest import testmod - - testmod() - print("Probability of 2 successes out of 4 trials") - print("with probability of 0.75 is:", end=" ") - print(binomial_distribution(2, 4, 0.75)) +"""For more information about the Binomial Distribution - +https://en.wikipedia.org/wiki/Binomial_distribution""" + +from math import comb + + +def binomial_distribution(successes: int, trials: int, prob: float) -> float: + """ + Return probability of k successes out of n tries, with p probability for one + success + + The function uses the factorial function in order to calculate the binomial + coefficient + + >>> binomial_distribution(3, 5, 0.7) + 0.30870000000000003 + >>> binomial_distribution (2, 4, 0.5) + 0.375 + >>> binomial_distribution(85, 171, 0.5) + 0.060749764626165485 + """ + if successes > trials: + raise ValueError("""successes must be lower or equal to trials""") + if trials < 0 or successes < 0: + raise ValueError("the function is defined for non-negative integers") + if not isinstance(successes, int) or not isinstance(trials, int): + raise ValueError("the function is defined for non-negative integers") + if not 0 < prob < 1: + raise ValueError("prob has to be in range of 1 - 0") + probability = (prob**successes) * ((1 - prob) ** (trials - successes)) + # Calculate the binomial coefficient: n! / k!(n-k)! + # math.comb is exact and avoids converting huge factorials to float. + coefficient = comb(trials, successes) + return probability * coefficient + + +if __name__ == "__main__": + from doctest import testmod + + testmod() + print("Probability of 2 successes out of 4 trials") + print("with probability of 0.75 is:", end=" ") + print(binomial_distribution(2, 4, 0.75))