MojoでProject Euler 61

https://projecteuler.net/problem=61

三角数から辿っていけばよいです。

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


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

fn collect_polygonals(p: Int, first: Int, last: Int) -> List[Int]:
    var ps = List[Int]()
    for n in range(1, last):
        var P = n * ((p - 2) * n - p + 4) // 2
        if P >= last:
            break
        if P >= first:
            ps.append(P)
    return ps

fn classify_by_last(ns: List[Int]) -> List[List[Int]]:
    var cl = initialize_list(100, List[Int]())
    for n in ns:
        var first = n[] % 100
        cl[first].append(n[])
    return cl

fn walk(v: Int, flag: Int, v0: Int, cl: List[List[List[Int]]]) -> Int:
    if flag == 63 << 3:
        if v // 100 == v0 % 100:
            return v
        else:
            return -1
    
    for p in range(3, 9):
        if ((flag >> p) & 1) != 0:
            continue
        var us = cl[p][v//100]
        for u in us:
            var s = walk(u[], flag | (1 << p), v0, cl)
            if s != -1:
                return v + s
    else:
        return -1

fn f() -> Int:
    var polygons = initialize_list(9, List[Int]())
    for p in range(3, 9):
        polygons[p] = collect_polygonals(p, 1000, 10000)
    
    var cl_last = initialize_list(9, List[List[Int]]())
    for p in range(3, 9):
        cl_last[p] = classify_by_last(polygons[p])
    
    var vs = polygons[3]
    for n in vs:
        var s = walk(n[], 1 << 3, n[], cl_last)
        if s != -1:
            return s
    else:
        return -1

fn main():
    print(f())