エラトステネスの篩→素因数分解→二項係数 with mod (JavaScript)

篩を使えば素因数分解できるし、素因数分解できるなら、そこそこ大きい二項係数 with mod も溢れず計算できるね、というだけの記事。

まず単純な実装のエラトステネスの篩。

const sieve = max => {
    const m = Math.ceil(Math.sqrt(max));
    const s = [];
    for (let p = 2; p <= m; p++)
        if (!s[p])
            for (let q = p * p; q <= max; q += p)
                s[q] = p;
    return s;
};

false / true ではなく、素因数のうちどれか一つ / 無ければ undefined が入っている(0 で初期化してもいいけどサボり)。
つまり、n が 2 以上で sieve[n] が undefined なら素数、と言うことになる。

こうしておくと次のようにして素因数分解ができる。

const factors = function* (n, s = sieve(n)) {
    if (n <= 1) return;
    for (; s[n]; n /= s[n]) yield s[n];
    yield n;
};
console.log([...factors(360)]); // [5, 3, 3, 2, 2, 2]

素因数分解ができたということは、競プロなんかでよく使うと噂の二項係数 with mod も書ける (特に速くはないけれど)。

とりあえず $ a^b \mod c $ を定義しまして...。

const modpow = (a, b, c) =>
    b < 1 ? 1 : b & 1 ?
        a * modpow(a, b - 1, c) % c :
        (modpow(a, b >> 1, c) ** 2) % c;

本題の $_nC_k\mod p$ はこんな感じ。

const nCkMp = (n, k, p, s = sieve(n)) => {
    k = Math.min(k, n - k);
    const a = [];
    for (let i = 0; i < k; i++) {
        for (const f of factors(n - i, s))
            a[f] = (a[f] | 0) + 1;
        for (const f of factors(k - i, s))
            a[f] = (a[f] | 0) - 1;
    }
    return a.reduce((n, v, k) => n * modpow(k % p, v, p) % p, 1);
};

$_nC_k$ が整数になることはわかりきってるので、素因数の多重集合 a で扱って重複度キャンセルして最後に乗算。

適当な値で測ってみる。

console.time();
let r = nCkMp(1000, 300, 1e9 + 7);
console.timeEnd();
console.log(r); // 626555557

手元の Node.js v12.13.1だと平均して 1.2 ms 前後。
割を食っているのは generator なんで、ここをインライン化すれば 0.7 ms くらい。
篩の実装頑張ればもう少し速くなるかも。
もっと値が大きい時は、篩を再利用したり foctors 内部をメモ化したりしないと悲しい感じになる気がする。

競プロ等だと逆元テーブル化とか使うだろうし出番はなさそうだけれど、法が素数に限定されない & 毎度変わっても大丈夫なのは割と便利かもしれない。

まあ、こんな方向でも求められるよね、とぬるいことを言ってお茶を濁しておしまい。