「题解」[UOJ 450]「集训队作业 2018」复读机

一台复读机的 EGF \(F(x)=\displaystyle\sum_{n\ge 0}[d|n]\frac{x^n}{n!}\)。

则有答案的 EGF \(G(x)=F^k(x)\)。我们要求的就是 \(\left[\dfrac{x^n}{n!}\right]G(x)\)。

• \(d=1\)

此时有

\[\begin{align*} F(x)&=\sum_{n\ge 0}\frac{x^n}{n!} \\ &=e^x \end{align*} \]

\[\begin{align*} G(x)&=F^k(x) \\ &=e^{kx} \end{align*} \]

则 \(\left[\dfrac{x^n}{n!}\right]G(x)=k^n\)。

时间复杂度 \(\Theta(\log n)\)。

• \(d=2\)

此时有

\[\begin{align*} F(x)&=1+\frac{x^2}{2!}+\frac{x^4}{4!}+\frac{x^6}{6!}+\cdots \\ &=\frac{1}{2}\left(\sum_{n\ge 0}\frac{x^n}{n!}+\sum_{n \ge 0}(-1)^n\frac{x^n}{n!}\right) \\ &=\frac{e^x+e^{-x}}{2} \end{align*} \]

\[\begin{align*} G(x)&=\frac{1}{2^k}\cdot\left(e^x+e^{-x}\right)^k \\ &=\frac{1}{2^k}\cdot\sum_{i=0}^{k}\binom{k}{i}e^ie^{-(k-i)x} \\ &=\frac{1}{2^k}\cdot\sum_{i=0}^{k}\binom{k}{i}e^{(2i-k)x} \end{align*} \]

则 \(\left[\dfrac{x^n}{n!}\right]G(x)=\dfrac{1}{2^k}\displaystyle\sum_{i=0}^{k}\binom{k}{i}(2i-k)^n\)。

时间复杂度 \(\Theta(k \log n)\)。

• \(d=3\)

由单位根反演可得（以下记 \(w=w_3\)）：

\[\begin{align*} F(x)&=\sum_{n \ge 0}[3|n]\frac{x^n}{n!} \\ &=\frac{1}{3}\sum_{n\ge 0}(w^{0n}+w^{1n}+w^{2n})\frac{x^n}{n!} \\ &=\frac{1}{3}\left(e^x+e^{wx}+e^{w^2x}\right) \end{align*} \]

用二项式定理扩展到多项式的形式进行展开得：

\[\begin{align*} G(x)&=\frac{1}{3^k}\cdot\left(e^x+e^{wx}+e^{w^2x}\right)^k \\ &=\frac{1}{3^k}\cdot\sum_{i=0}^{k}\sum_{j=0}^{k-i}\binom{k}{i,j,k-i-j}e^{ix}e^{jwx}e^{(k-i-j)w^2x} \\ &=\frac{1}{3^k}\cdot\sum_{i=0}^{k}\sum_{j=0}^{k-i}\binom{k}{i,j,k-i-j}e^{\left[i+jw+\left(k-i-j\right)w^2\right]x} \end{align*} \]

其中 \(\displaystyle\binom{k}{i,j,k-i-j}\) 为多重组合数。

则 \(\left[\dfrac{x^n}{n!}\right]G(x)=\dfrac{1}{3^k}\displaystyle\sum_{i=0}^{k}\sum_{j=0}^{k-i}\binom{k}{i,j,k-i-j}\left[i+jw+\left(k-i-j\right)w^2\right]^n\)。

时间复杂度 \(\Theta(k^2 \log n)\)。

代码

#include <bits/stdc++.h>

#define int unsigned long long

using namespace std;

const int N = 1e3, P = 19491001;

int n, k, d, ans, w, fac[N + 10], ifac[N + 10];

void InitC() {
    fac[0] = fac[1] = ifac[0] = ifac[1] = 1;
    for (int i = 2; i <= N; i++) ifac[i] = (P - P / i) * ifac[P % i] % P;
    for (int i = 2; i <= N; i++) fac[i] = fac[i - 1] * i % P, ifac[i] = ifac[i] * ifac[i - 1] % P;
}

int C(int x, int y) { return fac[x] * ifac[y] % P * ifac[x - y] % P; }
int C(int x, int y, int z) { return fac[x] * ifac[y] % P * ifac[z] % P * ifac[x - y - z] % P; }

int qpow(int a, int b) {
    int ret = 1;
    for (; b; b >>= 1, a = a * a % P) b & 1 ? ret = ret * a % P : 0;
    return ret;
}

struct OI {
    int RP, score;
} FJOI2022;

signed main() {
    FJOI2022.RP++, FJOI2022.score++;
    scanf("%llu%llu%llu", &n, &k, &d);
    if (d == 1) return printf("%llu\n", qpow(k, n)), 0;
    InitC();
    if (d == 2) {
        for (int i = 0; i <= k; i++) ans = (ans + C(k, i) * qpow(2 * i - k, n)) % P;
        return printf("%llu\n", ans * qpow(qpow(2, k), P - 2) % P), 0;
    }
    w = qpow(7, (P - 1) / 3);
    for (int i = 0; i <= k; i++)
        for (int j = 0; j <= k - i; j++) ans = (ans + C(k, i, j) * qpow((i + j * w + (k - i - j) * w * w) % P, n)) % P;
    return printf("%llu\n", ans * qpow(qpow(3, k), P - 2) % P), 0;
}