AtCoder Regular Contest 004: D - 表現の自由 ( Freedom of expression )

問題記録 AtCoder ARC ARC #004 DP 約数

https://atcoder.jp/contests/arc004/tasks/arc004_4

問題概要

整数N, Mが与えられる。NをM個の整数の積として表す方法が何通りあるか求めよ。構成要素が同じでも並び順が異なればそれぞれ1通りとして数えるものとする。

|N| <= 109

M <= 105

解法

掛け算の中に現れるのはNの約数だけなので、まずはこれを求めておく。その上で素朴な解答を考えると以下のようなDPが立つ。

dp(i, j): Nの約数をi回掛けていて、現在の積がjであるような場合の数

このDPでjとしてありうるのはNの約数だけである。なぜなら現在の積が一度Nの約数以外になってしまったら、それに何を掛けてももうNにはできないからである。よってNの約数の個数をKとかで表すことにすると、このDPはiの状態数がM, jの状態数がKであわせてO(MK)の状態数になる。遷移も愚直にNの約数全部で試すことにすると、DPの計算量は全体でO(MK2)になる。109までの数だと約数の個数は多くて1350個くらいになる（cf. 高度合成数）ので、これは全然間に合わない。

ここでMが大きい場合ほとんどの数は1か-1になるということに注目する。たとえば1の次に小さい数である2であっても、30回かそこら掛け続けると109を超えてしまう。すると残りの数は全部1か-1で埋めるしかなくなる。というわけで、まずは1と-1以外の数を使ってNを作る場合の数を求めてから、あとでそれらの残りの部分に1を埋めるやり方の総数を求める、という方法を使うことができる。

1と-1以外の数を使う場合の数を求めるパートは上述のDPをやればよい。ただし1と-1を掛けるような遷移は行わないことと、掛ける回数は最大でも30回程度でよくなったことに注意する。計算量はO(MK2)のMが30になったもので 30 * 13502 となり十分間に合う（負の約数を考慮すると4倍になるが、それでも108くらいなのでそれなりの速度の言語なら問題ないはず）。

あとは各dp(i, N), dp(i, -N)に対して、余った数のぶんだけ1か-1を詰め込んでいく方法が何通りあるかを数えればよい。どれだけ余っているかはiを見ればわかる。また符号を調整する必要があるが、n個の1に対して-1にする数を奇数個選ぶ場合の数の総和と偶数個選ぶ場合の数の総和はどちらも 2 ^ (n - 1) である（二項係数の足し算）。最後に並べ替え方であるが、n個の1とM-n個のそれ以外のもの（区別しない）を並べ替えるやりかたは MCn である。これらをdp(i, N)に掛けていけばよい。

感想

公式解説へのリンクがないのでこれが想定解法なのかわからない　DPの遷移はもっとうまく枝刈りできそうではあるけど

コード (D言語)

import std.stdio, std.array, std.string, std.conv, std.algorithm;
import std.typecons, std.range, std.random, std.math, std.container;
import std.numeric, std.bigint, core.bitop, std.bitmanip;

immutable long MOD = 10^^9 + 7;

void main() {
    auto s = readln.split.map!(to!long);
    auto N = s[0];
    auto M = s[1];

    if (N == 1) {
        if (M == 1) {
            writeln(1);
        } else {
            writeln(powmod(2, M-1, MOD));
        }
        return;
    }

    long[] F;
    for (long i = 1; i * i <= abs(N); ++i) {
        if (N % i == 0) {
            if (i * i == N) {
                F ~= i;
                F ~= -i;
            } else {
                F ~= i;
                F ~= N / i;
                F ~= -i;
                F ~= -(N / i);
            }
        }
    }

    F.sort!"abs(a) < abs(b)";

    int K = F.length.to!int;
    int[long] D;
    foreach (i, f; F) D[f] = i.to!int;

    const mx = 33;
    auto dp = new long[][](mx, K);
    dp[0][D[1]] = 1;

    foreach (i; 0..mx-1) {
        foreach (j; 0..K) {
            if (dp[i][j] == 0) continue;
            foreach (k; 2..K) {
                long x = F[j] * F[k];
                if (x !in D) continue;
                (dp[i+1][D[x]] += dp[i][j]) %= MOD;
            }
        }
    }

    long ans = 0;
    auto Comb = new Combination(10^^5+10);

    foreach (i; 0..min(mx, M+1)) {
        foreach (j; K-2..K) {
            if (dp[i][j] == 0) continue;
            if (i == M && F[j] == -N) continue;
            long n = M - i; // 1か-1でn個埋める
            if (n == 0) {
                ans += dp[i][j];
                ans %= MOD;
                continue;
            }
            long tmp = dp[i][j] * powmod(2, n - 1, MOD) % MOD;
            tmp = tmp * Comb.comb(M, n) % MOD;
            (ans += tmp) %= MOD;
        }
    }

    ans.writeln;
}


class Combination {
    long[] modinv;
    long[] f_mod;
    long[] f_modinv;

    this(int size) {
        modinv = new long[](size);
        modinv[0] = modinv[1] = 1;
        foreach(i; 2..size) {
            modinv[i] = modinv[MOD.to!int % i] * (MOD - MOD / i) % MOD;
        }

        f_mod = new long[](size);
        f_modinv = new long[](size);
        f_mod[0] = f_mod[1] = 1;
        f_modinv[0] = f_modinv[1] = 1;

        foreach(i; 2..size) {
            f_mod[i] = (i * f_mod[i-1]) % MOD;
            f_modinv[i] = (modinv[i] * f_modinv[i-1]) % MOD;
        }
    }

    long comb(long n, long k) {
        if (n < k) return 0;
        return f_mod[n] * f_modinv[n-k] % MOD * f_modinv[k] % MOD;
    }
}

long powmod(long a, long x, long m) {
    long ret = 1;
    while (x) {
        if (x % 2) ret = ret * a % m;
        a = a * a % m;
        x /= 2;
    }
    return ret;
}