MojoでProject Euler 90

https://projecteuler.net/problem=90

Stringは__getitem__で1文字を取れるんですね。

from math import sqrt
import sys


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

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 combinations_core[T: CollectionElement](a: List[T],
                                            n: Int, k: Int) -> List[List[T]]:
    if n == 0:
        return List[List[T]](List[T]())
    
    var vs = List[List[T]]()
    for k1 in range(k, len(a) - n + 1):
        var v = combinations_core(a, n-1, k1+1)
        for w in v:
            insert_list(w[], a[k1], 0)
            vs.append(w[])
    return vs

fn combinations[T: CollectionElement](a: List[T], n: Int) -> List[List[T]]:
    return combinations_core(a, n, 0)


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

fn exists_in_dice(d: Int, dice: List[Int]) -> Bool:
    return d in dice or (d == 6 and 9 in dice) or (d == 9 and 6 in dice)

fn is_valid_dies(dice1: List[Int], dice2: List[Int]) -> Bool:
    for n in range(1, 10):
        var d1 = n * n // 10
        var d2 = n * n % 10
        if ((not exists_in_dice(d1, dice1) or
                not exists_in_dice(d2, dice2)) and
                (not exists_in_dice(d1, dice2) or
                not exists_in_dice(d2, dice1))):
            return False
    return True

fn f() -> Int:
    var ds = List[Int]()
    for d in range(10):
        ds.append(d)
    var dies = combinations(ds, 6)
    
    var counter = 0
    for i in range(len(dies)):
        for j in range(i, len(dies)):
            if is_valid_dies(dies[i], dies[j]):
                counter += 1
    return counter

fn main():
    print(f())