MojoでProject Euler 46

https://projecteuler.net/problem=46

自然数nに対し、が素数なのか調べます。素数のSetを作っておけばいいのですが、答えがNなら、計算量はくらいなので、計算量をくらいに抑えたければ600000くらいまで素数を作っておけばよいです。

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


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

fn make_prime_table(N: Int) -> List[Int]:
    var a = initialize_list(N+1, True)
    for p in range(2, N+1):
        if p * p > N:
            break
        elif not a[p]:
            continue
        
        for n in range(p*p, N+1, p):
            a[n] = False
    
    var ps = List[Int]()
    for n in range(2, N+1):
        if a[n]:
            ps.append(n)
    return ps


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

fn f() -> Int:
    var N = 6*10**5
    var primes = make_prime_table(N)
    var set_primes = Set(primes)
    for n in range(9, N+1, 2):
        if n in primes:
            continue
        for m in range(1, n):
            var s = 2 * m * m
            if s >= n - 1:
                return n
            var p = n - s
            if p in set_primes:
                break
    else:
        return 0

fn main():
    print(f())