kmjp's blog

競技プログラミング参加記です

AtCoder ABC #241 : Ex - Card Deck Score

なるほど…。
https://atcoder.jp/contests/abc241/tasks/abc241_h

問題

N種類のカードがある。
i番目のカードは整数A[i]が書かれたものでB[i]ある。
これらsum(B)のカードからM枚を選び、書かれた整数の積を取る。
全てのカードの選び方における積の総和を求めよ。

解法

以下の多項式におけるM次の係数が解となる。
 \displaystyle f(x) = \prod(1+A_ix+ {A_i}^2x^2 \cdots + {A_i}^{B_i}x^{B_i})
等比数列の総和の式を使うと、
 \displaystyle f(x) = \prod(\frac{1-{A_i}^{B_i+1}x^{B_i+1}}{1-A_ix})
となる。

分母がN個の1次式の積なので、まずこれをN個の1次式の和に書き直そう。

分子についてはxの次数は高々2^N通りなので、それぞれの係数を求める。
そのそれぞれについて、上記N通りの1次式で割ったときのx^Mの係数を求めて行けばよい。

int N;
ll M,A[16],B[16],C[16];
const ll mo=998244353;

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


void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>M;
	FOR(i,N) cin>>A[i]>>B[i];
	
	FOR(i,N) {
		C[i]=1;
		FOR(j,N) if(j!=i) (C[i]*=(mo+1-A[j]*modpow(A[i])%mo))%=mo;
		C[i]=modpow(C[i]);
	}
	int mask;
	ll ret=0;
	FOR(mask,1<<N) {
		ll p=1,q=0;
		FOR(i,N) if(mask&(1<<i)) {
			p=(mo-p)*modpow(A[i],B[i]+1)%mo;
			q+=B[i]+1;
		}
		if(q<=M) {
			FOR(i,N) (ret+=C[i]*p%mo*modpow(A[i],M-q))%=mo;
		}
	}
	cout<<ret<<endl;
}

まとめ

最近多項式や母関数を使ったテク多いなぁ。