e092_2.py

from collections import Counter
from functools import reduce
from time import monotonic
from typing import Any, Generator, Sequence, TypeVar

T = TypeVar("T")


def step(n: int) -> int:
    return sum(i * i for i in map(int, str(n)))


def test(n: int) -> bool:
    while n != 89:
        if n == 1:
            return False
        n = step(n)
    return True


def comb_dup(xs: Sequence[T], n: int) -> Generator[Sequence[T], None, None]:
    def helper(
        acc: Sequence[T], tail: Sequence[T], k: int
    ) -> Generator[Sequence[T], None, None]:
        if k == 0:
            yield acc
            return
        for i, x in enumerate(tail):
            yield from helper((*acc, x), tail[i:], k - 1)

    yield from helper((), xs, n)


def factorial(n: int) -> int:
    if n < 2:
        return 1
    return reduce(lambda x, y: x * y, range(1, n + 1))


def calc(xs: Sequence[Any]) -> int:
    cnt = Counter(xs).values()
    s = factorial(sum(cnt))
    for t in cnt:
        s = s // factorial(t)
    return s


e = monotonic()
s: set[int] = {n for n in range(1, 568) if test(n)}
result: int = 0
for ns in comb_dup(range(10), 7):
    t = sum(i * i for i in ns)
    if t in s:
        result += calc(ns)
print(result)
print(f"{(monotonic() - e) * 1000:.2f}ms")