probability-calculator

About this project

This project is from freeCodeCamp’s Scientific Computing with Python Certificate. This repository contains the prompt for the project as well as my solution for the assignment.

To run my probability_calc code with the tests provided by freeCodeCamp for this project, feel free to visit my replit.

Assignment

Suppose there is a hat containing 5 blue balls, 4 red balls, and 2 green balls. What is the probability that a random draw of 4 balls will contain at least 1 red ball and 2 green balls? While it would be possible to calculate the probability using advanced mathematics, an easier way is to write a program to perform a large number of experiments to estimate an approximate probability.

For this project, you will write a program to determine the approximate probability of drawing certain balls randomly from a hat.

First, create a Hat class in prob_calculator.py. The class should take a variable number of arguments that specify the number of balls of each color that are in the hat. For example, a class object could be created in any of these ways:

hat1 = Hat(yellow=3, blue=2, green=6)
hat2 = Hat(red=5, orange=4)
hat3 = Hat(red=5, orange=4, black=1, blue=0, pink=2, striped=9)

A hat will always be created with at least one ball. The arguments passed into the hat object upon creation should be converted to a contents instance variable. contents should be a list of strings containing one item for each ball in the hat. Each item in the list should be a color name representing a single ball of that color. For example, if your hat is {"red": 2, "blue": 1}, contents should be ["red", "red", "blue"].

The Hat class should have a draw method that accepts an argument indicating the number of balls to draw from the hat. This method should remove balls at random from contents and return those balls as a list of strings. The balls should not go back into the hat during the draw, similar to an urn experiment without replacement. If the number of balls to draw exceeds the available quantity, return all the balls.

Next, create an experiment function in prob_calculator.py (not inside the Hat class). This function should accept the following arguments:

The experiment function should return a probability.

For example, let’s say that you want to determine the probability of getting at least 2 red balls and 1 green ball when you draw 5 balls from a hat containing 6 black, 4 red, and 3 green. To do this, we perform N experiments, count how many times M we get at least 2 red balls and 1 green ball, and estimate the probability as M/N. Each experiment consists of starting with a hat containing the specified balls, drawing a number of balls, and checking if we got the balls we were attempting to draw.

Here is how you would call the experiment function based on the example above with 2000 experiments:

hat = Hat(black=6, red=4, green=3)
probability = experiment(hat=hat, 
                  expected_balls={"red":2,"green":1},
                  num_balls_drawn=5,
                  num_experiments=2000)

Since this is based on random draws, the probability will be slightly different each time the code is run.

My solution

import copy
import random

class Hat:
    
    def __init__(self, **colors):
        self.d = {}
        self.contents = []

        for color in colors:
            self.d.update({color: colors[color]})

            for i in range(0, colors[color]):
                self.contents.append(color)

    def draw(self, number):
        if number > len(self.contents):
            self.balls_drawn = self.contents
            self.contents = []
            
            for color in self.d:
                self.d[color] = 0
        else:
            self.balls_drawn = random.sample(self.contents, k=number)

            for ball in self.balls_drawn:
                self.contents.remove(ball)
                self.d[ball] = self.d[ball] - 1

        return self.balls_drawn

def experiment(hat, expected_balls, num_balls_drawn, num_experiments):
    M = 0
    N = num_experiments

    for i in range(0, N):
        hat_copy = copy.deepcopy(hat)
        drawn = hat_copy.draw(num_balls_drawn)

        z = len(expected_balls)
        for color in expected_balls:
            if drawn.count(color) < expected_balls[color]:
                break
            else:
                z -= 1
        
        if z == 0:
            M += 1

    prob = M / N            
    return prob