MojoでProject Euler 92

https://projecteuler.net/problem=92

までとすると、次はまでになります。だから、そこまでの1か89のどちらにたどり着くかの表を作っておきます。そうしておいて重複組み合わせの一つ一つがどちらにたどり着くかを調べます。例えば、1111222と1212121は同じところにたどり着くのは明らかでしょう。1111222の並び替えは、

個あります。
こうすると、でも時間内に処理できます。

from math import sqrt
import sys


#################### List ####################

fn initialize_list[T: CollectionElement](N: Int, init: T) -> List[T]:
    var a = List[T](capacity=N)
    for n in range(N):
        a.append(init)
    return a

trait Printable(CollectionElement, Stringable):
    pass

fn print_list[T: Printable](a: List[T]):
    if a.size > 0:
        var s = "[" + str(a[0])
        for i in range(1, a.size):
            s += ", " + str(a[i])
        s += "]"
        print(s)
    else:
        print("[]")

fn duplicate_combinations(a: List[Int], n: Int) -> List[List[Int]]:
    var v = List[List[Int]]()
    if n == 1:
        for e in a:
            var w = List[Int](e[])
            v.append(w)
    elif n % 2 == 1:
        var v1 = duplicate_combinations(a, n-1)
        for e in a:
            for w1 in v1:
                if e[] > w1[][0]:
                    continue
                var w = List[Int](e[])
                w.extend(w1[])
                v.append(w)
    else:
        var m = n // 2
        var v1 = duplicate_combinations(a, m)
        for w1 in v1:
            for w2 in v1:
                if w1[][m-1] > w2[][0]:
                    continue
                var w = w1[]
                w.extend(w2[])
                v.append(w)
    return v

fn count_perm(a: List[Int]) -> Int:
    var L = len(a)
    var n = 1
    var m = 1
    var prev = a[0]
    for i in range(1, L):
        if a[i] != prev:
            for j in range(m):
                n *= L - i + j + 1
            for j in range(m):
                n //= j + 1
            m = 1
            prev = a[i]
        else:
            m += 1
    return n


#################### process ####################

fn sum_squares(owned n: Int) -> Int:
    var s = 0
    while n > 0:
        var d = n % 10
        n //= 10
        s += d**2
    return s

fn sum_squares_list(ds: List[Int]) -> Int:
    var s = 0
    for d in ds:
        s += d[]**2
    return s

fn make_destination_table(N: Int) -> List[Int]:
    var a = initialize_list(N+1, 0)
    for n in range(1, N+1):
        if a[n] != 0:
            continue
        var stack = List[Int](n)
        while True:
            n = sum_squares(n)
            if n == 1 or n == 89:
                break
            elif a[n] != 0:
                n = a[n]
                break
            stack.append(n)
        for e in stack:
            a[e[]] = n
    return a

fn f(E: Int) -> Int:
    var N = 10**E
    var table = make_destination_table(E*9**2)
    var ds = List[Int](0, 1, 2, 3, 4, 5, 6, 7, 8, 9)
    var counter = 0
    for ds1 in duplicate_combinations(ds, E):
        var s = sum_squares_list(ds1[])
        if table[s] == 89:
            counter += count_perm(ds1[])
    return counter

fn main() raises:
    var args = sys.argv()
    var E = atol(args[1])
    print(f(E))