#1192992 (C++17) No.3770 Sum of Pow of Distinct Count

• 問題: No.3770 Sum of Pow of Distinct Count
• ユーザー: marc2825
• 提出日時: 2026-10-13 03:38:00
• 言語: C++17 (gcc 15.3.0 + boost 1.92.0 + ACL)
• 結果: AC
• 実行時間: 563 ms / 4,000 ms
• コード長: 6,383 bytes
• コンパイル時間: 1,761 ms
• コンパイル使用メモリ: 226,760 KB
• 実行使用メモリ: 30,032 KB
• 最終ジャッジ日時: 2026-10-16 20:58:44
• 合計ジャッジ時間: 9,010 ms
• ジャッジサーバーID （参考情報）: judge2_0 / judge1_0

ファイルパターン	結果
sample	AC * 3
other	AC * 15

ソースコード

/// Claude Opus 5.5

#include <bits/stdc++.h>
using namespace std;

using u32 = uint32_t;
using u64 = uint64_t;

constexpr u32 MOD = 998244353;
constexpr u32 PRIMITIVE_ROOT = 3;

u64 mpow(u64 b, u64 e) {
    b %= MOD;
    u64 r = 1;
    while (e) {
        if (e & 1) r = r * b % MOD;
        b = b * b % MOD;
        e >>= 1;
    }
    return r;
}

// ---------------- NTT ----------------
vector<u32> rt, irt;  // rt[h + j] = ω_{2h}^j

void prepareRoots(int maxN) {
    rt.assign(maxN, 0);
    irt.assign(maxN, 0);
    for (int h = 1; h < maxN; h <<= 1) {
        u64 w = mpow(PRIMITIVE_ROOT, (MOD - 1) / (2 * h));
        u64 iw = mpow(w, MOD - 2);
        u64 cur = 1, icur = 1;
        for (int j = 0; j < h; j++) {
            rt[h + j] = (u32)cur;
            irt[h + j] = (u32)icur;
            cur = cur * w % MOD;
            icur = icur * iw % MOD;
        }
    }
}

void ntt(vector<u32>& a, bool inverse) {
    const int n = (int)a.size();
    for (int i = 1, j = 0; i < n; i++) {
        int bit = n >> 1;
        for (; j & bit; bit >>= 1) j ^= bit;
        j ^= bit;
        if (i < j) swap(a[i], a[j]);
    }
    const vector<u32>& R = inverse ? irt : rt;
    for (int h = 1; h < n; h <<= 1) {
        for (int i = 0; i < n; i += 2 * h) {
            for (int j = 0; j < h; j++) {
                u32 x = a[i + j];
                u32 y = (u32)((u64)a[i + j + h] * R[h + j] % MOD);
                u32 s = x + y;
                if (s >= MOD) s -= MOD;
                u32 d = x + MOD - y;
                if (d >= MOD) d -= MOD;
                a[i + j] = s;
                a[i + j + h] = d;
            }
        }
    }
    if (inverse) {
        const u64 ninv = mpow(n, MOD - 2);
        for (auto& v : a) v = (u32)(v * ninv % MOD);
    }
}

vector<u32> multiply(vector<u32> a, vector<u32> b) {
    if (a.empty() || b.empty()) return {};
    const int need = (int)(a.size() + b.size() - 1);
    if (min(a.size(), b.size()) <= 32) {
        vector<u64> c(need, 0);
        for (size_t i = 0; i < a.size(); i++)
            for (size_t j = 0; j < b.size(); j++)
                c[i + j] = (c[i + j] + (u64)a[i] * b[j]) % MOD;
        return vector<u32>(c.begin(), c.end());
    }
    int n = 1;
    while (n < need) n <<= 1;
    a.resize(n);
    b.resize(n);
    ntt(a, false);
    ntt(b, false);
    for (int i = 0; i < n; i++) a[i] = (u32)((u64)a[i] * b[i] % MOD);
    ntt(a, true);
    a.resize(need);
    return a;
}

// f * g ≡ 1 (mod x^m), f[0] != 0
vector<u32> inverseSeries(const vector<u32>& f, int m) {
    vector<u32> g = {(u32)mpow(f[0], MOD - 2)};
    int cur = 1;
    while (cur < m) {
        const int nxt = cur * 2, sz = nxt * 2;
        vector<u32> fa(sz, 0), ga(sz, 0);
        for (int i = 0; i < min((int)f.size(), nxt); i++) fa[i] = f[i];
        for (int i = 0; i < cur; i++) ga[i] = g[i];
        ntt(fa, false);
        ntt(ga, false);
        for (int i = 0; i < sz; i++) {
            u64 fg = (u64)fa[i] * ga[i] % MOD;
            fa[i] = (u32)((2 + MOD - fg) % MOD * ga[i] % MOD);
        }
        ntt(fa, true);
        fa.resize(nxt);
        g.swap(fa);
        cur = nxt;
    }
    g.resize(m);
    return g;
}

// Σ_{l∈[lo,hi)} c[l] / (1 - x[l]·t) = P / Q   (|P| = len, |Q| = len + 1)
void sumOfFractions(int lo, int hi, const vector<u32>& c, const vector<u32>& x,
                    vector<u32>& P, vector<u32>& Q) {
    const int len = hi - lo;
    if (len <= 32) {
        P.assign(len, 0);
        Q.assign(len + 1, 0);
        Q[0] = 1;
        for (int l = lo, d = 0; l < hi; l++, d++) {
            const u64 cl = c[l];
            const u64 nx = (MOD - x[l]) % MOD;  // -x[l]
            // P <- P·(1 - x t) + c·Q
            for (int i = d; i >= 0; i--) {
                u64 v = P[i];
                if (i > 0) v += nx * P[i - 1] % MOD;
                v += cl * Q[i] % MOD;
                P[i] = (u32)(v % MOD);
            }
            // Q <- Q·(1 - x t)
            for (int i = d + 1; i >= 1; i--) {
                Q[i] = (u32)((Q[i] + nx * Q[i - 1]) % MOD);
            }
        }
        return;
    }
    const int mid = (lo + hi) / 2;
    vector<u32> PL, QL, PR, QR;
    sumOfFractions(lo, mid, c, x, PL, QL);
    sumOfFractions(mid, hi, c, x, PR, QR);
    int n = 1;
    while (n < len + 1) n <<= 1;
    PL.resize(n);
    QL.resize(n);
    PR.resize(n);
    QR.resize(n);
    ntt(PL, false);
    ntt(QL, false);
    ntt(PR, false);
    ntt(QR, false);
    P.assign(n, 0);
    Q.assign(n, 0);
    for (int i = 0; i < n; i++) {
        P[i] = (u32)(((u64)PL[i] * QR[i] + (u64)PR[i] * QL[i]) % MOD);
        Q[i] = (u32)((u64)QL[i] * QR[i] % MOD);
    }
    ntt(P, true);
    ntt(Q, true);
    P.resize(len);
    Q.resize(len + 1);
}

int main() {
    u64 N, M;
    int K;
    if (!(cin >> N >> M >> K)) return 0;

    prepareRoots(1 << 20);

    vector<u64> fact(K + 1), ifact(K + 1);
    fact[0] = 1;
    for (int i = 1; i <= K; i++) fact[i] = fact[i - 1] * i % MOD;
    ifact[K] = mpow(fact[K], MOD - 2);
    for (int i = K; i >= 1; i--) ifact[i - 1] = ifact[i] * i % MOD;

    const u64 Mm = M % MOD;
    const u64 Ne = N % (MOD - 1);

    // w_i = C(M, i) · (M - i)^N
    vector<u32> w(K + 1);
    u64 fall = 1;  // M (M-1) ... (M-i+1) mod p
    for (int i = 0; i <= K; i++) {
        const u64 base = (Mm + MOD - i) % MOD;
        const u64 binom = fall * ifact[i] % MOD;
        const u64 pw = (base == 0) ? 0 : mpow(base, Ne);
        w[i] = (u32)(binom * pw % MOD);
        fall = fall * base % MOD;
    }

    // c_l = Σ_{i=l}^{K} (-1)^{i-l} C(i, l) w_i
    vector<u32> a(K + 1), b(K + 1);
    for (int j = 0; j <= K; j++) {
        a[j] = (u32)((u64)w[K - j] * fact[K - j] % MOD);
        b[j] = (j & 1) ? (u32)((MOD - ifact[j]) % MOD) : (u32)ifact[j];
    }
    const vector<u32> conv = multiply(a, b);
    vector<u32> c(K + 1), x(K + 1);
    for (int l = 0; l <= K; l++) {
        c[l] = (u32)((u64)conv[K - l] * ifact[l] % MOD);
        x[l] = (u32)((Mm + MOD - l) % MOD);
    }

    // Σ_k S_k t^k = Σ_l c_l / (1 - (M - l) t)
    vector<u32> P, Q;
    sumOfFractions(0, K + 1, c, x, P, Q);
    vector<u32> S = multiply(P, inverseSeries(Q, K + 1));
    S.resize(K + 1);

    string out;
    out.reserve((size_t)K * 11);
    for (int k = 1; k <= K; k++) {
        out += to_string(S[k]);
        out += '\n';
    }
    fwrite(out.data(), 1, out.size(), stdout);
    return 0;
}