题解 [HAOI2018] 染色

题目链接

Luogu
Loj

题目描述

对一个长度为 \(n\) 的序列染色，颜色一共有 \(m\) 种，对于每一种染色方案，设有 \(u\) 种颜色恰好出现了 \(s\) 次，则将答案加 \(w_u\) ，求答案
其中 \(n \leq 10^7, m \leq 10^5, s \leq 150\)，并将答案对 \(1004535809\) 取模

分析

下面的式子虽然有点多，但每一步推导都非常简单，应该可以在不使用草稿纸的情况下看懂

首先，\(1004535809\) 的出现直接暗示了我们这题需要 NTT

然后肯定就是生成函数推柿子题了

考虑枚举恰好出现了 \(s\) 次的颜色有多少种

就假设枚举到 \(i\) 种

那么对于这 \(i\) 种颜色，他们的指数型生成函数都是

\[\frac{x^s}{s!} \]

那对于另外的 \(m - i\) 种颜色，由于他们不能恰好出现 \(s\) 次，那他们的生成函数就应该是

\[\begin{equation} \begin{split} F(x) &= \frac{x^0}{0!} + \frac{x^1}{1!} + \frac{x^2}{2!} + \frac{x^3}{3!} + ... - \frac{x^s}{s!}\\ &= \sum_{k \ge 0} \frac{x^k}{k!} - \frac{x^s}{s!}\\ &= e^x - \frac{x^s}{s!} \end{split} \end{equation} \]

那么，假设选定了 \(i\) 种颜色使他们恰好出现 \(s\) 次，这样的序列个数就是

\[n! \times [x^n](\frac{x^s}{s!})^i (e^x - \frac{x^s}{s!})^{m - i} \]

这个 \(n!\) 是因为这是指数型生成函数

再考虑从 \(m\) 种颜色中选出 \(i\) 种颜色作为恰好出现 \(s\) 次的颜色的方案数，所以答案就是

\[\begin{equation} \begin{split} ans &= \sum_{i = 0}^m w_i \times \binom{m}{i} \times n! \times [x^n](\frac{x^s}{s!})^i (e^x - \frac{x^s}{s!})^{m - i}\\ &= \sum_{i = 0}^m w_i \times \binom{m}{i} \times \frac{n!}{(s!)^i} \times [x^n]x^{si} (e^x - \frac{x^s}{s!})^{m - i}\\ &= \sum_{i = 0}^m w_i \times \binom{m}{i} \times \frac{n!}{(s!)^i} \times [x^{n - si}] (e^x - \frac{x^s}{s!})^{m - i}\\ \end{split} \end{equation} \]

对于最后那一项，我们用二项式定理展开，有

\[(e^x - \frac{x^s}{s!})^{m - i} = \sum_{j = 0}^{m - i} (-1)^j \frac{x^{sj}}{(s!)^j} e^{x(m-i-j)} \binom{m-i}{j} \]

代入得答案为

\[ans = \sum_{i = 0}^m w_i \times \binom{m}{i} \times \frac{n!}{(s!)^i} \times [x^{n - si}] \sum_{j = 0}^{m - i} (-1)^j \frac{x^{sj}}{(s!)^j} e^{x(m-i-j)} \binom{m-i}{j}\\ \]

这里观察到如果 \(j\) 从 \(i\) 开始枚举的话式子会简单许多

\[ans = \sum_{i = 0}^m w_i \times \binom{m}{i} \times \frac{n!}{(s!)^i} \times [x^{n - si}] \sum_{j = i}^{m} (-1)^{j-i} \frac{x^{s(j-i)}}{(s!)^{j-i}} e^{x(m-j)} \binom{m-i}{j-i} \]

把 \(\frac{1}{(s!)^i}\) 和 \([x^{n - si}]\) 写进去

\[\begin{equation} \begin{split} ans &= \sum_{i = 0}^m w_i \times \binom{m}{i} \times n! \times \sum_{j = i}^{m} [x^{n - sj}] (-1)^{j-i} \frac{1}{(s!)^j} e^{x(m-j)} \binom{m-i}{j-i}\\ &=\sum_{i = 0}^m w_i \times \binom{m}{i} \times n! \times \sum_{j = i}^{m} (-1)^{j-i} \frac{1}{(s!)^j} \binom{m-i}{j-i} [x^{n - sj}] e^{x(m-j)} \end{split} \end{equation} \]

而我们又知道

\[e^{kx} = \sum_{i \ge 0} \frac{k^i s^i}{i!} \]

所以

\[[x^{n - sj}] e^{x(m-j)} = \frac{(m-j)^{n-sj}}{(n-sj)!} \]

代进去

\[ans = \sum_{i = 0}^m w_i \times \binom{m}{i} \times n! \times \sum_{j = i}^{m} (-1)^{j-i} \frac{1}{(s!)^j} \binom{m-i}{j-i} \frac{(m-j)^{n-sj}}{(n-sj)!} \]

观察发现两个二项式系数显然可以消掉一个阶乘，于是我们把二项式系数展开

\[ans = \sum_{i = 0}^m w_i \times \frac{m!}{i!} \times n! \times \sum_{j = i}^{m} (-1)^{j-i} \frac{1}{(s!)^j} \frac{1}{(j-i)!(m-j)!} \frac{(m-j)^{n-sj}}{(n-sj)!} \]

其实到这里已经差不多可以看出来了，我们再变形一下

\[ans = \sum_{i = 0}^m w_i \times \frac{m!}{i!} \times n! \times \sum_{j = i}^{m} \frac{(-1)^{j-i}}{(j-i)!} \times \frac{(m-j)^{n-sj}}{(s!)^j(m-j)!(n-sj)!} \]

可以发现 \(\frac{(-1)^{j-i}}{(j-i)!}\) 的取值只与 \(j-i\) 有关，\(\frac{(m-j)^{n-sj}}{(s!)^j(m-j)!(n-sj)!}\) 只与 \(j\) 有关
我们分别将其设为 \(g_{j-i}\) 与 \(f_j\)

那么

\[ans = \sum_{i = 0}^m w_i \times \frac{m!}{i!} \times n! \times \sum_{j = i}^{m} f_j \times g_{j-i} \]

可以发现这里有一个极其类似卷积的式子，对于这样的式子，我们的操作是将 \(f\) 翻转，即让 \(f_j\) 与 \(f_{m-j}\) 交换，那么对于新的 \(f_j\)，他对应的原来的 \(f\) 值就是 \(f_{m-j}\)，所以他应该乘的 \(g\) 值就是 \(g_{m-j-i}\)，到这里卷积的形式就出来了

我们再写一遍答案

\[\begin{equation} \begin{split} ans &= \sum_{i = 0}^m w_i \times \frac{m!}{i!} \times n! \times \sum_{j = 0}^{m-i} f_j \times g_{m-j-i}\\ &= n!m! \sum_{i = 0}^m \frac{w_i}{i!} \sum_{j = 0}^{m-i} f_j \times g_{m-j-i} \end{split} \end{equation} \]

其中

\[f_i = \frac{(-1)^{m-i}}{(m-i)!}\\ g_i = \frac{(m-i)^{n-si}}{(s!)^i(m-i)!(n-si)!} \]

如果 \(n-si < 0\) 则令 \(f_i = 0\) 即可

使用 NTT 求解卷积，复杂度为 \(O(n + m\log m)\)

代码

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

const int N = 1e7 + 10;
const int M = 400010;
const LL p = 1004535809;

LL n, m, s, ge = 3, ig, mx = 1, ans, w[M], f[M], g[M], fac[N];

LL ksm(LL v, LL tms)
{
    LL res = 1;
    for(; tms; tms >>= 1, v = v * v % p) if(tms & 1) res = res * v % p;
    return res;
}

void NTT(LL *a, LL mx, LL flag)
{
    for(int i = 0, j = 0; i < mx; i++){
        if(i < j) swap(a[i], a[j]);
        for(int l = mx >> 1; (j ^= l) < l; l >>= 1);
    }
    for(int i = 2; i <= mx; i <<= 1){
        LL now = (i >> 1), step = ksm((flag > 0 ? ge : ig), (p - 1) / i);
        for(int j = 0; j < mx; j += i){
            LL temp = 1;
            for(int k = 0; k < now; k++){
                LL G = a[j + k], H = temp * a[j + k + now] % p;
                a[j + k] = (G + H) % p, a[j + k + now] = (G - H + p) % p;
                temp = step * temp % p;
            }
        }
    }
    LL inv = ksm(mx, p - 2);
    if(flag < 0) for(int i = 0; i < mx; i++) a[i] = a[i] * inv % p;
}

int main()
{
    ios :: sync_with_stdio(false), cin.tie(0);
    ig = ksm(ge, p - 2);
    cin >> n >> m >> s;
    for(int i = 0; i <= m; i++) cin >> w[i];

    fac[0] = 1;
    for(int i = 1; i <= max(n, m); i++) fac[i] = fac[i - 1] * i % p;

    for(int i = 0; i <= m; i++)
        g[i] = (((i & 1) ? -1 : 1) * ksm(fac[i], p - 2) + p) % p;
    for(int i = 0; i <= m; i++)
        if(s * i > n) f[i] = 0; //特判f[i]=0
        else f[i] = ksm((ksm(fac[s], i) * fac[m - i] % p) * fac[n - s * i] % p, p - 2) * ksm(m - i, n - s * i) % p;
    for(int i = 0; (i << 1) <= m; i++) swap(f[i], f[m - i]); //翻转f

    while(mx <= m) mx <<= 1; mx <<= 1;
    NTT(f, mx, 1);
    NTT(g, mx, 1);
    for(int i = 0; i < mx; i++) f[i] = f[i] * g[i] % p;
    NTT(f, mx, -1);

    for(int i = 0; i <= m; i++)
        if(s * i <= n) ans = (ans + ((w[i] * ksm(fac[i], p - 2) % p) * f[m - i] % p) + p) % p;
    ans = ans * (fac[n] * fac[m] % p) % p;
    printf("%lld\n", ans);
    return 0;
}