from random import randint
from math import factorial
from collections import defaultdict
from fractions import Fraction


def binomial(n: int, k: int) -> Fraction:
    return Fraction(factorial(n), (factorial(k) * factorial(n - k)))


def f(k: int, n: int) -> Fraction:
    # Probability that k out of n trials succeed with rate p per trial.
    p = Fraction(1, 2)
    return binomial(n, k) * p**k * (1 - p) ** (n - k)


def count(exp_n: Fraction, step: int, dist) -> Fraction:
    step += 1

    # For each number `n` of remaining "ones" calculate the probability that
    # 0..`n` ones are remanining for that probability. We can then calculate a
    # new probability mass by multiplying the current with the new probability
    # for a given `n`.
    ndist = defaultdict(Fraction)
    for n, p in dist:
        for i in range(n + 1):
            np = f(i, n)
            ndist[n - i] += p * np

    # Probability that we reach `0` at this step counts to total expected
    # number of steps.
    exp_n += ndist[0] * step
    del ndist[0]

    # Drop probabilities small enough that they seem englible for the precision
    # we need which is 10 digits after the decimal point.
    dist = []
    for n, p in ndist.items():
        if p > Fraction(1, 10**15):
            dist.append((n, p))
    if dist:
        return count(exp_n, step, dist)
    return exp_n


def stackoverflow():
    # I found this expression when googling (geminiing) an expression for this
    # type of problem. However, it did not yield the correct solution because I
    # did not use Fractions. I then came up with `count` and was somehow clever
    # enough to realize that I would have to use Fractions and after I came
    # back I updated stackoverflow accordingly just to see if I had implemented
    # the expression correctly.

    # https://mathoverflow.net/questions/41604/what-is-the-expected-maximum-out-of-a-sample-size-n-from-a-geometric-distribut
    t = Fraction()
    p = Fraction(1, 2)
    q = 1 - p
    n = 32
    for i in range(1, n + 1):
        c = binomial(n, i)
        c *= (-1) ** (i + 1)
        c *= Fraction(1, 1 - q**i)
        t += c
    return float(t)


def euler_323():
    n_stackoverflow = stackoverflow()
    n = count(Fraction(), 0, [(32, Fraction(1, 1))])
    assert round(float(n_stackoverflow), 10) == round(float(n), 10)
    return round(float(n), 10)


if __name__ == "__main__":
    solution = euler_323()
    print(f"e323.py: {solution}")
    assert solution == 6.3551758451