AtCoder ABC #235 : G - Gardens

ちょっとあてずっぽうでやってしまった。
https://atcoder.jp/contests/abc235/tasks/abc235_g

問題

N個の庭がある。
ここに3種類の苗がそれぞれA・B・C個あり、それぞれ庭に植えることを考える。

• 1つの庭に、同種の苗は1つまでしか植えられない
• 苗が余ってもよい

この時、N個の庭にそれぞれ1個以上苗が植えてあるような苗の植え方は何通りか。

解法

g(m) := 苗をあるm個の庭に植える植え方。ただし、1個も苗のない庭があってもよい
とすると、包除原理の要領で sum(g(m)*binom(N,m)*(-1)^(N-m) が解となる。

f(m,l) := m個の庭に、l番目の種類の苗を植える植え方。余った苗があってもよい。
とすると、g(m)=f(m,1)*f(m,2)*f(m,3)となる。
あとはf(m,l)を植える植え方を考えよう。

• mが苗の個数以下なら、各庭にその苗を植えても植えなくても良いのでf(m,l) = 2^mである。
• mが苗の個数より多い場合、例えばその苗の個数をPとするとf(m,l)=binom(m,0)+binom(m,1)+...+binom(m,P)となる。

後者は逐一二項係数を(P+1)解計算すると間に合わない。
ただ、f(m+1,k)=2*f(m,l)-binom(m,P)であることを用いると、二項係数のあらかじめ前処理でO(1)で求められるようにしておけば、f(m,l)をそれぞれO(N)で埋めきることができる。

int N,A[3];
const ll mo=998244353;
ll AP[3][5050505];

const int NUM_=10400001;
static ll fact[NUM_+1],factr[NUM_+1],inv[NUM_+1];

ll modpow(ll a, ll n = mo-2) {
	ll r=1;a%=mo;
	while(n) r=r*((n%2)?a:1)%mo,a=a*a%mo,n>>=1;
	return r;
}

ll comb(ll N_, ll C_) {
	if (fact[0]==0) {
		inv[1]=fact[0]=factr[0]=1;
		for (int i=2;i<=NUM_;++i) inv[i] = inv[mo % i] * (mo - mo / i) % mo;
		for (int i=1;i<=NUM_;++i) fact[i]=fact[i-1]*i%mo, factr[i]=factr[i-1]*inv[i]%mo;
	}
	if(C_<0 || C_>N_) return 0;
	return factr[C_]*fact[N_]%mo*factr[N_-C_]%mo;
}
ll hcomb(ll P_,ll Q_) { return (P_==0&&Q_==0)?1:comb(P_+Q_-1,Q_);}

ll G[5050505];

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>A[0]>>A[1]>>A[2];
	
	FOR(i,3) {
		for(j=0;j<=N;j++) {
			if(A[i]>=j) {
				AP[i][j]=modpow(2,j);
			}
			else {
				AP[i][j]=(AP[i][j-1]*2+mo-comb(j-1,A[i]))%mo;
			}
		}
	}
	
	ll ret=0;
	for(i=0;i<=N;i++) {
		G[i]=AP[0][i]*AP[1][i]%mo*AP[2][i]%mo;
		if((N-i)%2==0) ret+=G[i]*comb(N,i)%mo;
		else ret+=mo-G[i]*comb(N,i)%mo;
	}
	cout<<ret%mo<<endl;
}

まとめ

二項係数の和を高速に求めたいとき、パスカルの三角形を思い浮かべると高速な計算方法を思いつきやすい気がする。