MojoでProject Euler 77

https://projecteuler.net/problem=77

五角数定理みたいなものはないので、素直に

を計算します。

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

trait Printable(CollectionElement, Stringable):
    pass

fn print_list[T: Printable](a: List[T]):
    if a.size > 0:
        var s = "[" + str(a[0])
        for i in range(1, a.size):
            s += ", " + str(a[i])
        s += "]"
        print(s)
    else:
        print("[]")


#################### 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 multiple(p: Int, inout a: List[Int]):
    for i in range(len(a) - p - 1, -1, -1):
        var j = i + p
        a[j] -= a[i]

fn inverse(a: List[Int]) -> List[Int]:
    var N = len(a) - 1
    var b = initialize_list(N+1, 0)
    b[0] = 1
    var c = List[Int]()
    for i in range(N+1):
        var e = b[i]
        c.append(e)
        for j in range(N+1-i):
            b[i+j] -= e * a[j]
    return c

fn f_each(N: Int) -> Int:
    var primes = make_prime_table(N)
    var a = initialize_list(N+1, 0)
    a[0] = 1
    for p in primes:
        multiple(p[], a)
    var c = inverse(a)
    return c[N]

fn F(N: Int) -> Int:
    var first = 1
    while f_each(first) < N:
        first *= 2
    first //= 2
    var last = first * 2
    while last - first > 2:
        var mid = (first + last) // 2
        if f_each(mid) > N:
            last = mid + 1
        else:
            first = mid
    if f_each(first) > N:
        return first
    else:
        return first + 1

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