MojoでProject Euler 74

https://projecteuler.net/problem=74

辿っていくとρ型になるのですが、すでにあと何歩進めるか分かっているノードに突き当たったら、そこで終了できます。

import sys


#################### library ####################

fn make_factorial_table(N: Int) -> List[Int]:
    var a = List[Int]()
    a.append(1)
    for n in range(1, N+1):
        a.append(a[n-1] * n)
    return a

fn digits(owned n: Int) -> List[Int]:
    var ds = List[Int]()
    while n > 0:
        var d = n % 10
        n //= 10
        ds.append(d)
    return ds

fn frequency[T: KeyElement](v: List[T]) -> Dict[T, Int]:
    var freq = Dict[T, Int]()
    for e in v:
        var f = freq.get(e[], 0)
        freq[e[]] = f + 1
    return freq


#################### 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("[]")


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

var facts = List[Int]()

fn step(n: Int) -> Int:
    var ds = digits(n)
    var s = 0
    for d in ds:
        s += facts[d[]]
    return s

fn find_rho(n: Int, a: List[Int], dists: List[Int]) -> List[Int]:
    var rho = List[Int](n)
    var p = a[n]
    var q = a[p]    # 2歩ずつ進む
    rho.append(p)
    while p != q:
        if dists[p] != 0:   # すでに確定したノードに突き当たる
            return rho
        p = a[p]
        q = a[a[q]]
        rho.append(p)
    rho.append(a[p])
    return rho

fn update_rho(rho: List[Int], inout dists: List[Int]):
    var L = len(rho)
    var pt = rho[L-1]
    # loopの始点はどこか？
    var k = 0
    while rho[k] != pt:
        k += 1
    
    for i in range(L-1):
        if i < k:
            dists[rho[i]] = L - i
        else:
            dists[rho[i]] = L - k - 1

fn update_I(rho: List[Int], inout dists: List[Int]):
    var L = len(rho)
    var pt = rho[L-1]
    # loopの始点はどこか？
    var k = 0
    while rho[k] != pt:
        if dists[rho[k]] != 0:
            break
        k += 1
    
    for i in range(k):
        dists[rho[i]] = k - i + dists[rho[k]]

fn dist_naive(owned n: Int) -> Int:
    var s = Set[Int]()
    while n not in s:
        s.add(n)
        n = step(n)
    return len(s)

fn max_number(N: Int) -> Int:
    if N < 10**7:
        return max(N, 2+6*facts[9])
    else:
        return N

fn f(N: Int) -> Int:
    facts = make_factorial_table(9)
    var M = max_number(N)
    var a = initialize_list(M+1, 0)
    for n in range(M+1):
        a[n] = step(n)
    
    var dists = initialize_list(M+1, 0)
    for n in range(1, M+1):
        if dists[n] != 0:
            continue
        var rho = find_rho(n, a, dists)
        if dists[rho[len(rho)-1]] == 0:     # ρ型
            update_rho(rho, dists)
        else:                               # I型
            update_I(rho, dists)
    
    var counter = 0
    for n in range(N+1):
        if dists[n] == 61:
            counter += 1
    return counter

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