题解-CodeChef IOPC14L Sweets Problem

Problem

CodeChef-IOPC14L

题目概要：给定 $n$ 种糖果且给定每种糖果的数量 $A_i$，$Q$ 组询问，每次问选出 $S$ 个糖果的方案数（模$10^9+7$）

$n\leq 10^6,A_i\leq 10^3,Q\leq 10^4,S\leq 2\times 10^3$

Solution

都说这题是容斥，但是始终不知道如何容斥，下面介绍一个母函数的做法

这题想暴力首先可以想到将所有糖果的母函数乘起来。形式化的，对于一种糖果若有 $t$ 个，则其母函数为 $\sum_{i=0}^tx^i$，将所有共 $n$ 个母函数乘起来得到的$x^S$对应的系数即为答案

问题转化为求所有母函数的积，设为 $F$

考虑到 $n\leq 10^6$，暴力求积复杂度为 $O(nS\log S)$，但$A_i\leq 10^3$，则考虑可以将本质相同的糖果合并一下，可以得出 $s[x]$ 表示有 $x$ 颗糖果的种类有多少，则

\[F=\prod_{i=1}^{2000}(\sum_{j=0}^ix^j)^{s[i]}
\]

这个东西的复杂度是 $O(S^2\log S\log n)$ 的，照理说已经比较优了，但还可以继续优化

考虑到

\[\sum_{i=0}^nx^i=\frac {x^{n+1}-1}{x-1}
\]

所以上面的式子可以化为

\[F=\prod_{i=1}^{2000}(\frac {x^{i+1}-1}{x-1})^{s[i]}
\]

\[F=\frac {\prod_{i=1}^{2000}(x^{i+1}-1)^{s[i]}}{(x-1)^n}
\]

上式中的分子由于每一项都是二项式的幂次，最终的答案可以利用Dp在 $O(AS\ln S)$ 的时间内求得（其中$\ln S$为级数求和复杂度）

至于分母，视作 $(x-1)^{-n}$，可以考虑

\[(1-x)^{-1}=\frac 1{1-x}=\sum_{i=0}^{+\infty}x^i
\]

\[(x-1)^{-n}=(-1)^n(1-x)^{-n}=(-1)^n(\frac 1{1-x})^n=(-1)^n(\sum_{i=0}^{+\infty}x^i)^n
\]

发现那个 $(-1)^n$ 很讨厌，于是将原式 $F$ 的分母分子同时乘一个 $(-1)^n$，由于 $\sum s[i]=n$，所以分子变为 $\prod_{i=1}^{2000}(1-x^{i+1})^{s[i]}$，同样可以利用Dp在 $O(AS\ln S)$ 的时间内求得

继续考虑分母，现在已经求得分母倒数为 $(\sum_{i=0}^{+\infty}x^i)^n$，这个式子中 $x^k$ 的项系数等价于将 $k$ 分解为 $n$ 个非负整数和的方案数，即$\binom {n+k-1}{k}$

统计答案只需要分子与分母倒数做卷积即可，时间复杂度 $O(n+AS\ln S+S^2)$

Code

另外，这道题不支持c++提交，整的我交了好多次才知道 c 与 c++ 的不同（比如说不能用取址符、不能用 const 定义出的数字去定义数组大小等）

#include <stdio.h>

#include <ctype.h>

typedef long long ll;

inline int read(){

	char c11=getchar();int x = 0;

	while(!isdigit(c11))c11=getchar();

	while(isdigit(c11))x=x*10+c11-'0',c11=getchar();

	return x;

}

//const int N = 2001, M = 2001000, p = 1000000007;

const int p = 1000000007;

#define N 2001

#define M 2001000

int fac[M],inv[M],s[N];

int f[N][N],Ans[N];

int n;

inline int qpow(int A,int B){

	int res = 1;while(B){

		if(B&1) res = (ll)res * A%p;

		A = (ll)A * A%p, B >>= 1;

	}return res;

}

inline int c(int nn,int mm){return (ll)fac[nn]*inv[mm]%p*inv[nn-mm]%p;}

int main() {

	fac[0] = f[0][0] = 1;

	for(int i=1;i<M;++i) fac[i] = (ll)fac[i-1]*i%p;

	inv[M-1] = qpow(fac[M-1],p-2);

	for(int i=M-1;i;--i) inv[i-1] = (ll)inv[i]*i%p;

	n = read();

	for(int i=1;i<=n;++i)++s[read()];

	for(int i=1;i<N;++i)

		for(int j=0,t;j<=s[i] && (t = (i+1)*j)<N;++j)

			if(j&1) for(int k=0;k+t<N;++k)

				f[i][k+t] = (f[i][k+t] + (ll)c(s[i],j) * f[i-1][k]%p*(p-1))%p;

			else for(int k=0;k+t<N;++k)

				f[i][k+t] = (f[i][k+t] + (ll)c(s[i],j) * f[i-1][k])%p;

	for(int i=0;i<N;++i)

		s[i] = c(n+i-1,i);

	for(int i=0;i<N;++i)

		for(int j=0;j<=i;++j)

			Ans[i] = (Ans[i] + (ll)f[N-1][j] * s[i-j])%p;

	int Q = read();

	while(Q--) printf("%d\n",Ans[read()]);

	return 0;

}