538-riddler/2020_08_08/main.py

from fractions import Fraction as F
from collections import defaultdict
from itertools import product
from math import factorial


class NematodeNumber:

    P = F(1, 2)

    def __init__(self, n):
        if isinstance(n, int):
            self._val = defaultdict(lambda: F(0))
            self._val[n] = F(1)
        else:
            self._val = n

    @classmethod
    def _breed(cls, n):

        def binom(n, k):
            n_choose_k = factorial(n) // (factorial(k) * factorial(n - k))
            return n_choose_k * cls.P**k * (1 - cls.P)**(n - k)

        leftover, pairs = n % 2, n // 2
        result = defaultdict(lambda: F(0))
        for n_bab in range(pairs + 1):
            prob = binom(pairs, n_bab)
            result[pairs*2 + n_bab + leftover] = prob
        return result

    def __add__(self, other):
        left = self._val
        right = other._val
        result = defaultdict(lambda: F(0))
        for (l, lp), (r, rp) in product(left.items(), right.items()):
            for b, bp in self._breed(l+r).items():
                result[b] += lp*rp*bp
        return NematodeNumber(result)

    def __pow__(self, exponent):
        base = self._val
        for _ in range(exponent):
            next_ = defaultdict(lambda: F(0))
            for n, np in base.items():
                for b, bp in self._breed(n).items():
                    next_[b] += np*bp
            base = next_
        return NematodeNumber(base)

    def __repr__(self):
        return repr(self._val)

    def expectation_value(self):
        return sum(val*prob for (val, prob) in self._val.items())

def limit():
    one = NematodeNumber(1)
    base_result = one + one
    for i in range(1, 30):
        res = base_result ** i
        print(f"{i: 3}: {res.expectation_value()}")
    """
      1: 3.0
      2: 3.625
      3: 4.40625
      4: 5.3828125 -> 5 49/128
      5: 6.603515625
      6: 8.12939453125
      7: 10.0367431640625
      8: 12.420928955078125
      9: 15.401161193847656
     10: 19.12645149230957
     11: 23.783064365386963
     12: 29.603830456733704
     13: 36.87978807091713
     14: 45.97473508864641
     15: 57.343418860808015
     16: 71.55427357601002
     17: 89.31784197001252
    """


def main():
    # r1 = add({2: F(1)}, {2: F(1)})
    # print(r1)
    # r2 = add(r1, r1)
    # print(r2)
    # print(exp({2: F(1)}, 2))
    # val = exp(add({1: F(1)}, {1: F(1)}), 4)
    limit()
    # print(float(expectation_value(val)))

if __name__ == "__main__":
    main()