64 lines
1.6 KiB
Python
64 lines
1.6 KiB
Python
from functools import lru_cache
|
|
|
|
|
|
BALLS_PER_COLOR = 10
|
|
N_COLORS = 7
|
|
TO_DRAW = 20
|
|
|
|
|
|
@lru_cache(maxsize=10**9)
|
|
def count(urn, to_draw):
|
|
if to_draw == 0:
|
|
distinct_colors = sum([1 if c != BALLS_PER_COLOR else 0 for c in urn])
|
|
return distinct_colors
|
|
|
|
remaining_balls = sum(urn)
|
|
expected = 0
|
|
for color_i, color_remaining in enumerate(urn):
|
|
if color_remaining == 0:
|
|
continue
|
|
new_urn = list(urn)
|
|
new_urn[color_i] -= 1
|
|
expected_for_color = count(tuple(sorted(new_urn)), to_draw - 1)
|
|
probability = color_remaining / remaining_balls
|
|
expected += (probability * expected_for_color)
|
|
return expected
|
|
|
|
|
|
def binomial_coefficient(n, k):
|
|
if k < 0 or k > n:
|
|
return 0
|
|
if k == 0 or k == n:
|
|
return 1
|
|
k = min(k, n - k)
|
|
result = 1
|
|
for i in range(k):
|
|
result *= n - i
|
|
result //= i + 1
|
|
return result
|
|
|
|
|
|
def math():
|
|
# Find probability that a given color is chose in TO_DRAW draws.
|
|
combs_without_color = binomial_coefficient((N_COLORS - 1) * BALLS_PER_COLOR, TO_DRAW)
|
|
combs_with_color = binomial_coefficient(N_COLORS * BALLS_PER_COLOR, TO_DRAW)
|
|
p_single_color = 1 - (combs_without_color / combs_with_color)
|
|
# Then, expected colors is that probability times the number of colors.
|
|
return p_single_color * N_COLORS
|
|
|
|
|
|
def euler_493():
|
|
urn = [BALLS_PER_COLOR for _ in range(N_COLORS)]
|
|
c = count(tuple(urn), TO_DRAW)
|
|
c = round(c, 9)
|
|
m = round(math(), 9)
|
|
assert c == m
|
|
return c
|
|
|
|
|
|
if __name__ == "__main__":
|
|
solution = euler_493()
|
|
print("e493.py: " + str(solution))
|
|
# assert(solution == 0)
|
|
|