Solution -「NOI Simu.」逆天题

$\mathscr{Description}$

  对于 $r=0,1,\cdots,n-1$, 设 $\{1,2,\cdots,nm\}$ 中有 $f_r$ 个子集满足子集内元素之和 $\bmod n=r$, 求出

$$\sum_{r=0}^{n-1}f_r^2\bmod 998244353.$$

  $n,m\le10^{18}$.

$\mathscr{Solution}$

  我们知道 $f$ 的 GF:

$$F(x)=\left(\prod_{i=0}^{n-1}(1+x^i)\right)^m\bmod (x^n-1).$$

于是我们需要用 $\omega_n$ 对 $F$ 做 DFT. 不妨设 $\hat{f_k}=\operatorname{DFT}(\boldsymbol f)_k$, 那么

$$\hat{f_k}=\left(\prod_{i=0}^{n-1}(1+\omega_n^{ik})\right)^m.$$

  看过 3B1B 的视频, 我们已知了对乘积式 $\prod_{i=0}^{n-1}(1+\omega_n^i)$ 的处理技巧:

  注意到 $1-z^n=0\Leftrightarrow z=\omega_n^k$, 那么 $1-z^n=\prod_{i=0}^{n-1}(z-\omega_n^i)$. 代入 $z=-1$, 观察左右, 我们有 $\prod_{i=0}^{n-1}(1+\omega_n^i)=2[2\nmid n]$.

指标乘上常数 $k$, 无非让 $\omega_n$ 转圈圈的速度发生了一些变化. 令 $d=\gcd(k,n)$, 容易得到

$$\hat{f_k}=[2\nmid (n/d)]2^{dm}.$$

  答案应当为 $\boldsymbol f^2$, 也即是 $\sum_{i=0}^{n-1}f_i^2$. 注意到 $\gcd(k,n)=\gcd(n-k,n)$, 于是 $f_i=f_{n-i}$, 所以只需要求

$$\sum_{i=0}^{n-1}f_if_{-i\bmod n}=[x^0]F^2(x).$$

如果能算 $\hat{\boldsymbol f}$, 我们就能立马得到 $\operatorname{DFT}(F^2)$, 接下来只需要提取 IDFT 矩阵的第一行 (带上 $\frac{1}{n}$ 的系数) 与 $F^2$ 的 DFT 向量做点乘就能得到答案.

  如何算 $\hat{\boldsymbol f}$? 将 $n$ 素因数分解, 暴力枚举 $d$ 算贡献即可. 使用 Pollard-Rho, 动态维护 $\varphi(n/d)$, 预处理 $2$ 的光速幂, 可以做到 $\mathcal O(\sqrt P+Td(n))$, 其中 $P=998244353$, $d(n)$ 为 $n$ 的因子个数 (其自然比 Pollard-Rho 的 $n^{1/4}$ 高阶).

$\mathscr{Code}$

/*+Rainybunny+*/

// #include <bits/stdc++.h>
#include <bits/extc++.h>

#define rep(i, l, r) for (int i = l, rep##i = r; i <= rep##i; ++i)
#define per(i, r, l) for (int i = r, per##i = l; i >= per##i; --i)

typedef long long LL;
#define fi first
#define se second

const int MOD = 998244353;
LL n, m;
__gnu_pbds::gp_hash_table<LL, int> fct;

namespace Factor {

inline LL add(LL u, const LL v, const LL m) {
    return (u += v) < m ? u : u - m;
}

inline LL mul(const LL u, const LL v, const LL m) {
    return __int128(u) * v % m;
}

inline LL mpow(LL u, LL v, LL m) {
    LL ret = 1;
    for (; v; u = mul(u, u, m), v >>= 1) ret = mul(ret, v & 1 ? u : 1, m);
    return ret;
}

inline LL gcd(const LL u, const LL v) { return v ? gcd(v, u % v) : u; }

inline LL labs(const LL u) { return u < 0 ? -u : u; }

inline bool millerRabin(const LL x, const LL b) {
    LL k = x - 1; while (!(k & 1)) k >>= 1;
    static LL pwr[70]; pwr[0] = 1, pwr[1] = mpow(b, k, x);
    while (k != x - 1) {
        pwr[pwr[0] + 1] = mul(pwr[pwr[0]], pwr[pwr[0]], x);
        ++pwr[0], k <<= 1;
    }
    per (i, pwr[0], 1) {
        if (pwr[i] != 1 && pwr[i] != x - 1) return false;
        if (pwr[i] == x - 1) return true;
    }
    return true;
}

inline bool isprime(const LL x) {
    if (x == 2 || x == 3 || x == 5 || x == 7 || x == 11) return true;
    if (x == 61 || x == 127) return true;
    if (!(x % 2) || !(x % 3) || !(x % 5) || !(x % 7) || !(x % 11))return false;
    if (!(x % 61) || !(x % 127)) return false;
    return millerRabin(x, 2) && millerRabin(x, 61) && millerRabin(x, 127);
}

inline LL pollardRho(const LL x) {
    static std::mt19937 emt(time(0) ^ 20120712);
    for (LL a = emt() % (x - 1) + 1, len = 1, st = 0, ed = 0; ;
      len <<= 1, st = ed) {
        LL prd = 1;
        rep (stp, 1, len) {
            prd = mul(prd, labs(st - (ed = add(mul(ed, ed, x), a, x))), x);
            if (!(stp & 127) && gcd(prd, x) > 1) return gcd(prd, x);
        }
        if (gcd(prd, x) > 1) return gcd(prd, x);
    }
}

inline void factor(LL x, __gnu_pbds::gp_hash_table<LL, int>& res,
  const int k = 1) {
    if (x == 1) return ;
    if (isprime(x)) return void(res[x] += k);
    LL d = pollardRho(x); int cnt = 0;
    while (!(x % d)) x /= d, ++cnt;
    factor(x, res, k), factor(d, res, k * cnt);
}

} // namespace Factor

inline int mul(const int u, const int v) { return 1ll * u * v % MOD; }
inline void addeq(int& u, const int v) { (u += v) >= MOD && (u -= MOD); }
inline int mpow(int u, int v) {
    int ret = 1;
    for (; v; u = mul(u, u), v >>= 1) ret = mul(ret, v & 1 ? u : 1);
    return ret;
}

namespace Power2 {

const int SM = 31596;
int sma[SM + 1], gnt[SM + 1];

inline void init() {
    sma[0] = 1;
    rep (i, 1, SM) addeq(sma[i] = sma[i - 1], sma[i - 1]);
    gnt[0] = 1, gnt[1] = sma[SM];
    rep (i, 2, SM) gnt[i] = mul(gnt[i - 1], gnt[1]);
}

inline int power(const int u) {
    return mul(sma[u % SM], gnt[u / SM]);
}

} // namespace Power2

inline int calc(const __gnu_pbds::gp_hash_table<LL, int>
  ::iterator&& it, LL d, LL phi) {
    if (it == fct.end()) {
        return mul(phi % MOD, Power2::power(2 * d % (MOD - 1)
          * (m % (MOD - 1)) % (MOD - 1)));
    }

    int ret = 0;
    rep (i, 0, it->se) {
        if (i == repi || it->fi != 2) addeq(ret, calc(std::next(it), d, phi));
        if (i == repi) break;
        d *= it->fi, phi /= it->fi - (i + 1 == repi);
    }
    return ret;
}

int main() {
    freopen("ntt.in", "r", stdin);
    freopen("ntt.out", "w", stdout);

    Power2::init();
    int T; scanf("%d", &T);
    while (T--) {
        scanf("%lld %lld", &n, &m);
        fct.clear(), Factor::factor(n, fct);
        LL phi = n;
        for (const auto& p: fct) phi = phi / p.fi * (p.fi - 1);
        printf("%d\n", mul(mpow(n % MOD, MOD - 2), calc(fct.begin(), 1, phi)));
    }
    return 0;
}