MojoでProject Euler 27

https://projecteuler.net/problem=27

素数の判定をどれくらい大きさまで行っていいか分からないので、ナイーブに判定してメモ化してみました。
MojoではDictを使えるようになりました。

from collections import Dict

fn main() raises:
    var d = Dict[String, Int]()
    d["a"] = 1
    print(d["a"])       # 1
    print("a" in d)     # True

なので、これをメモ化に使おうとしましたが、

from collections import Dict

var memo = Dict[String, Int]()
fn f(s: String) raises -> Int:
    if s in memo:
        return memo[s]
    elif s == "a":
        memo[s] = 1
        return 1
    else:
        memo[s] = 2
        return 2

fn main() raises:
    print(f("a"))
    print(f("b"))
    print(f("a"))
    print(f("b"))

このテストコードはなぜか落ちます。どうも現状Dictをグローバルな変数にはできないっぽいです。仕方なく、DynamicVectorを使いました。

from collections import Dict
from math import min
import sys


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

def div_pow(n: Int, d: Int) -> Tuple[Int, Int]:
    var m = n
    var e = 0
    while m % d == 0:
        e += 1
        m //= d
    return (e, m)


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

fn calc_min_b(a: Int, L: Int) -> Int:
    if a > 0:
        return 2
    elif -a < L * 2:
        var n = a // 2
        return 2 - (n * n + a * n)
    else:
        return 2 - (L * L + a * L)

var primes = DynamicVector[Int]()
def is_prime(n: Int) -> Bool:
    if n < 2:
        return False
    elif n < primes.size and primes[n] != -1:
        return primes[n] == 1
    
    var b = True
    for p in range(2, n+1):
        if p * p > n:
            break
        elif n % p == 0:
            b = False
            break
    
    if n >= primes.size:
        for m in range(primes.size, n+1):
            primes.push_back(-1)
    primes[n] = 1 if b else 0
    return b

def length(a: Int, b: Int) -> Int:
    var n = 0
    while True:
        m = n * n + a * n + b
        if not is_prime(m):
            break
        n += 1
    return n

fn f(N: Int) raises -> Int:
    var max_length = 0
    var max_product = 0
    for a in range(-N+1, N):
        var min_b = calc_min_b(a, max_length)
        for b in range(min_b, N+1):
            var l = length(a, b)
            if l > max_length:
                max_length = l
                max_product = a * b
                print(a, b, l)
    return max_product

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