MojoでProject Euler 87

https://projecteuler.net/problem=87

同じ値で違う組み合わせがあるので、ナイーブに組むしかなさそうです。

import sys


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

fn int_sqrt(n: Int) -> Int:
    var x = 1
    x = (x + n // x) // 2
    while True:
        var x1 = (x + n // x) // 2
        if x1 >= x:
            return x
        x = x1


#################### 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


#################### prime ####################

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(N: Int) -> Int:
    var primes = make_prime_table(int_sqrt(N))
    var s = Set[Int]()
    for p2 in primes:
        var n2 = p2[]**2
        for p3 in primes:
            var n3 = p3[]**3
            if n2 + n3 >= N:
                break
            for p4 in primes:
                var n4 = p4[]**4
                var n = n2 + n3 + n4
                if n >= N:
                    break
                s.add(n)
    return len(s)

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