46 lines
1.4 KiB
Python
46 lines
1.4 KiB
Python
import copy
|
|
import random
|
|
from collections import defaultdict
|
|
|
|
# Consider using the modules imported above.
|
|
|
|
class Hat:
|
|
def __init__(self, **kwargs: dict[str, int]) -> None:
|
|
self.contents: list[str] = []
|
|
for color, count in kwargs.items():
|
|
[self.contents.append(color) for _ in range(count)]
|
|
|
|
def __repr__(self) -> str:
|
|
return str(self.contents)
|
|
|
|
def draw(self, n_balls: int):
|
|
drawn = []
|
|
if (n_balls >= len(self.contents)):
|
|
drawn = self.contents
|
|
self.contents = []
|
|
else:
|
|
for _ in range(n_balls):
|
|
index = random.randrange(len(self.contents))
|
|
drawn.append(self.contents[index])
|
|
self.contents[index] = self.contents[-1]
|
|
self.contents.pop()
|
|
return drawn
|
|
|
|
|
|
def experiment(hat: Hat, expected_balls: dict[str, int], num_balls_drawn: int, num_experiments: int):
|
|
n_desired_result = 0
|
|
for _ in range(num_experiments):
|
|
hat_copy = copy.deepcopy(hat)
|
|
drawn = hat_copy.draw(num_balls_drawn)
|
|
drawn_dict = defaultdict(lambda: 0)
|
|
for ball in drawn:
|
|
drawn_dict[ball] += 1
|
|
desired_outcome = True
|
|
for ball in expected_balls:
|
|
if (drawn_dict[ball] < expected_balls[ball]):
|
|
desired_outcome = False
|
|
break
|
|
if desired_outcome:
|
|
n_desired_result += 1
|
|
return n_desired_result/num_experiments
|