MojoでProject Euler 31（2）

https://projecteuler.net/problem=31

この問題は母関数が使えます。具体的にはNにする方法の数を求めるとすると、
$\displaystyle G(x) = \prod_p{\sum_{\substack{p|n \\ 0 \le n \le N}}{x^n}}$
ここでpは1, 2, 5, ..., 200を取ります。もっと具体的に書くと、
$G(x) = (1 + x + x^2 + \cdots + x^{200})(1 + x^2 + x^4 + \cdots + x^{200})\cdots$
例えば、3にする方法は、1を3枚か1と2を1枚ずつですが、
母関数をこう書くと、
$G(x) = (x^{1 \times 0}+x^{1 \times 1}+x^{1 \times 2}+\cdots)(x^{2 \times 0}+x^{2 \times 1}+\cdots)\cdots$
上の2つの方法は、 $x^{1 \times 3}$ と $x^{1 \times 1}x^{2 \times 1}$ に対応することが分かります。
上の母関数は、 $x^{200}$ より次数の大きな項を書いても結果に影響を与えないので、
$\displaystyle G(x) = \prod_p{\sum_{\substack{p|n \\ 0 \le n}}{x^n}}$
これは、
$\displaystyle G(x) = \prod_p{\frac{1}{1-x^p}}$
とシンプルに書けます。
要するに割り算をしていくということですが、この割り算はO(N)で可能で、pは $O(\log{N})$ 個あるから全体の計算量は $O(N\log{N})$ になります。

import sys


#################### constatns ####################

alias D = 10**9+7


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

fn collect_coins(N: Int) -> List[Int]:
    var ns = List[Int](1, 2, 5);
    var ps = List[Int]()
    for e in range(19):
        for n in ns:
            var p = n[] * 10**e
            if p > N:
                break
            ps.append(p)
    return ps

# g / (1 - x^e)
fn divide(g: List[Int], e: Int) -> List[Int]:
    var h = g
    for i in range(e, g.size):
        h[i] = (h[i] + h[i-e]) % D
    return h

fn f(N: Int) -> Int:
    var ps = collect_coins(N)
    var g = List[Int]()
    g.append(1)
    for i in range(1, N+1):
        g.append(0)
    for p in ps:
        g = divide(g, p[])
    return g[N]

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