[ABC285Ex] Avoid Square Number 【生成函数】【容斥原理】

AysctLucky / 2024-02-04 / 原文

很好的题目，使我螺旋升天。

讨论题目要求，容易想到容斥，设 $x_i$ 表示钦定有 $i$ 个元素是完全平方数的方案数，则有答案：

\[\sum_{i=0}^{n}\begin{pmatrix}n \\ i\end{pmatrix}(-1)^{i}x_i \]

考虑如何求 $x_i$，显然我们要把每个 $E_j$ 分成至少 $i$ 个偶数。

设这个方案数为 $f_{i,E_j}$，发现它可以很方便的用生成函数表达：

\[G_i(x) = (1-x^2)^{-i} \times (1-x)^{i - n} = (1+x)^{-i}(1-x)^{-n} \]

考虑到 $(1-x)^{-1}$ 相当于前缀和，$(1+x)^{-1}$ 相当于前缀差，我们可以在 $O(nw)$ 的时间求出 $G_0$，并通过 $O(w)$ 的时间从 $G_{i-1}$ 转移到 $G_i$，所以总时间复杂度为 $O(nw)$（$w = \max_j(e_j)$）。

在求出 $G_i$ 后，我们可以通过 $O(k)$ 的时间复杂度求出 $x_i$。

总时间复杂度 $O(nw+nk)$。

#include<bits/stdc++.h>
#define int long long
#define poly vector<int>
const int mod = 1e9 + 7;
using namespace std;
int fpow(int a,int b)
{
    int r = 1;
    while(b)
    {
        if(b & 1)r = r * a % mod;
        a = a * a % mod;
        b >>= 1;
    }
    return r % mod;
}
void presum(poly& f)
{
    for(int i=1,sz=f.size();i<sz;i++)
    {
        f[i] = (f[i] + f[i-1]) % mod;
    }
}
void presub(poly& f)
{
    for(int i=1,sz=f.size();i<sz;i++)
    {
        f[i] = (f[i] - f[i-1] + mod) % mod;
    }
}
int n,k,e[10005],mxe;
int c[10005];
int ans;
poly f;
signed main()
{
    ios::sync_with_stdio(false);
    cin.tie(0);
    cout.tie(0);
    cin>>n>>k;
    for(int i=1;i<=k;i++)
    {
        cin>>e[i];
        mxe = max(mxe,e[i]);
    }
    f.resize(mxe+1);
    f[0] = 1;
    for(int i=1;i<=n;i++)presum(f);
    c[0] = 1;
    for(int i=1;i<=n;i++)
    {
        c[i] = c[i-1] * (n - i + 1) % mod * fpow(i,mod-2) % mod;
    }
    for(int i=0;i<=n;i++)
    {
        int res = c[i];
        if(i & 1)res = mod - res;
        for(int j=1;j<=k;j++)res = res * f[e[j]] % mod;
        ans = (ans + res) % mod;
        presub(f);
    }
    cout<<ans;
    return 0;
}