MojoでProject Euler 73

https://projecteuler.net/problem=73

しらみつぶしをしてもよいのですが、もっと速い方法があります。
まず、約分を考えなければ、dを固定して考えて、分母の上限をNと書くと、

 \displaystyle \sum_d{\lfloor \frac{N-1}{2} \rfloor} - \sum_d{\lfloor \frac{N}{3} \rfloor}

となります。これをがんばって計算すると、

 \lfloor \frac{(N-1)^2}{4} \rfloor - \lfloor \frac{N(N-1)}{6} \rfloor

そして、約分できない分数の個数を S(N)と書くと、分母と分の最大公約数が2となるものが S(\lfloor \frac{N}{2} \rfloor)、などとすれば、

 \displaystyle S(N) = \lfloor \frac{(N-1)^2}{4} \rfloor - \lfloor \frac{N(N-1)}{6} \rfloor - \sum_{d=2}^{\lfloor \frac{N}{5} \rfloor}{S(\lfloor \frac{N}{d} \rfloor)}

となります。さらに、dが大きいところでは同じ値になるので、まとめて計算できます。
こうすれば、 O(N)よりだいぶ速くなります。ただ、すぐにオーバーフローするので、 N = 2 \times 10^9くらいが限度です(上の数式を見て分かるように、少し工夫すればもうちょっと大きなNでも計算できます)。

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("[]")


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

var memo = Dict[Int, Int]()
fn f(N: Int) -> Int:
    if N <= 4:
        return 0
    var m = memo.get(N, -1)
    if m != -1:
        return m
    
    var s = (N - 1)**2 // 4 - N * (N - 1) // 6
    var k = 5
    while True:
        var w = N // k - N // (k + 1)
        if w <= 1:
            break
        s -= f(k) * w
        k += 1
    for d in range(2, N // k + 1):
        s -= f(N // d)
    memo[N] = s
    return s

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