题解-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;
}