Codeforces Hello 2019: D. Makoto and a Blackboard

https://codeforces.com/contest/1097/problem/D

問題概要

整数N, Kが与えられる。Nの約数のうちひとつを等確率で選び、それでNを割る、という操作をK回繰り返すとき、最終的なNの値の期待値を求めよ。

N <= 1015

K <= 104

解法

Nの約数の数は最大ケースで30000個近くあるため、約数列挙して素朴なDPをやるのでは間に合わない。

Nを素因数分解した上で操作を考えてみると、一度の操作は素因数分解の指数のところをランダムに減らしていく操作に相当する。ここで各素因数について指数の減り方は独立なので、素因数ごとに独立に問題を解いていくことができる。つまりE(n, k)を整数nに対してk回操作を行ったあとの期待値とすると、

N = P0 ^ a × P1 ^ b × P2 ^ c × ...

素因数分解されるとき、

E(N, K) = E(P0 ^ a, K) × E(P1 ^ b, K) × E(P2 ^ c, K) × ...

という風に独立にできる。指数部は最大でも50程度の値までしか取らず、素因数の数も同程度に収まるので、各素因数ごとにDPをしてEを求めていっても十分間に合う。

注意点として元のNが大きいせいでMODを取らずに掛け算するとオーバーフローする場合があるので気をつける。

感想

素因数ごとに独立!素因数ごとに独立!素因数ごとに独立!素因数ごとに独立!素因数ごとに独立!素因数ごとに独立!

コード (C++)

#include <bits/stdc++.h>
using namespace std;
#define REP(i,n) for (int i=0;i<(n);i++)
#define REP2(i,m,n) for (int i=m;i<(n);i++)
typedef long long ll;
typedef long double ld;
typedef vector<int> VI;
typedef vector<ll> VL;

const ll MOD = 1000000007;

ll N, K, M;
vector<ll> P;
vector<ll> C;

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

ll solve2(ll p, ll c) {
    vector<VL> dp = vector<VL>(K+1, VL(c+3, 0));
    dp[0][c] = 1;

    REP(i, K) {
        REP(j, c+1) {
            ll v = dp[i][j] * powmod(j+1, MOD-2, MOD) % MOD;
            (dp[i+1][0] += v) %= MOD;
            (dp[i+1][j+1] -= v) %= MOD;
        }
        REP(j, c+1) {
            (dp[i+1][j+1] += dp[i+1][j]) %= MOD;
        }
    }

    ll ans = 0;

    REP(j, c+1) {
        ans += dp[K][j] * powmod(p % MOD, j, MOD) % MOD;
        ans %= MOD;
    }

    return (ans % MOD + MOD) % MOD;
}

void solve() {
    cin >> N >> K;
    ll n = N;

    for (ll i = 2; i * i <= n; ++i) {
        if (n % i == 0) {
            P.push_back(i);
            C.push_back(0);
            M += 1;
        }
        while (n % i == 0) {
            C.back() += 1;
            n /= i;
        }
    }

    if (n > 1) {
        P.push_back(n);
        C.push_back(1);
        M += 1;
    }

    ll ans = 1;

    REP(i, M) {
        ans = ans * solve2(P[i] % MOD, C[i]) % MOD;
    }

    cout << ans << endl;
}

int main() {
    cin.tie(0);
    ios::sync_with_stdio(false);

    solve();
}