MojoでProject Euler 95

https://projecteuler.net/problem=95

辿っていって二度同じノードが出てきたらそこでループしていると判定して、ループ上の場合は周期をメモして、そうでないときは-1を記録していきます。

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 range_list(first: Int, last: Int) -> List[Int]:
    var a = List[Int]()
    for n in range(first, last):
        a.append(n)
    return a

fn insert_list[T: CollectionElement](inout v: List[T], e: T, i: Int):
    v.append(e)
    for k in range(len(v) - 1, i, -1):
        v[k] = v[k-1]
    v[i] = e

fn max(a: List[Int]) -> Int:
    var max_value = a[0]
    for i in range(1, len(a)):
        if a[i] > max_value:
            max_value = a[i]
    return max_value


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

fn make_sum_divisors_table(N: Int) -> List[Int]:
    var a = range_list(0, N+1)
    var b = initialize_list(N+1, 1)
    
    for p in range(2, N+1):
        if p*p > N:
            break
        if a[p] != p:
            continue
        for n in range(p, N+1, p):
            var m = 1
            while a[n] % p == 0:
                a[n] //= p
                m = m * p + 1
            b[n] *= m
    
    for n in range(2, N+1):
        if a[n] != 1:
            b[n] *= a[n] + 1
    
    for n in range(2, N+1):
        b[n] -= n
    return b

fn walk(v0: Int, inout a: List[Int], table: List[Int]):
    var N = len(a) - 1
    var visited = Set[Int]()
    visited.add(v0)
    var path = List[Int](v0)
    var v = v0
    while True:
        v = table[v]
        if v > N or a[v] == -1 or v == 1:
            for v1 in path:
                a[v1[]] = -1
            return
        elif v in visited:  # 前に通った
            var L = len(path)
            var loop_len = 0
            for i in range(L):
                var v1 = path[i]
                if v1 == v:
                    loop_len = L - i
                if loop_len == 0:   # まだループ上じゃない
                    a[v1] = -1
                else:
                    a[v1] = loop_len
            return
        else:
            path.append(v)
            visited.add(v)

fn f(N: Int) -> Int:
    var table = make_sum_divisors_table(N)
    var a = initialize_list(N+1, 0)
    for n in range(2, N+1):
        walk(n, a, table)
    
    var max_loop_len = max(a)
    for v in range(2, N+1):
        if a[v] == max_loop_len:
            return v
    return -1

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