kmjp's blog

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

AtCoder ARC #106 : D - Powers

Codeforcesにありそう。
https://atcoder.jp/contests/arc106/tasks/arc106_d

問題

N要素の整数列Aが与えられる。
入力Kが与えられるので、1≦X≦Kとなる各整数Xに対し、 \left(\displaystyle \sum_{L=1}^{N-1} \sum_{R=L+1}^{N} (A_L+A_R)^X\right) \bmod 998244353を答えよ。

解法

 \displaystyle 2 \times \sum_{L=1}^{N-1} \sum_{R=L+1}^{N} (A_L+A_R)^X + \sum_{L=1}^{N} (A_L+A_L)^X = \sum_{L=1}^{N} \sum_{R=1}^{N} (A_L+A_R)^X
であることを利用する。左辺の1項目を直接求めるとO(N^2*X)かかりそうで面倒だが、2項目はO(N*X)でいける。
右辺もLとRが対称であることを利用して解こう。

 \displaystyle (A_L+A_R)^X = \sum_i C(X,i)\times{A_L}^i\times{A_R}^{X-i}であることを利用すると、
 \displaystyle \sum_L \sum_R (A_L+A_R)^X = \sum_i C(X,i)\times\sum_L{A_L}^i\times\sum_R{A_R}^{X-i}となる。
A[i]のX乗和を事前に計算しておけば、この総和は高速に計算可能である。

int N,K;
ll A[202020][305];
ll S[305];
const ll mo=998244353;

ll comb(ll N_, ll C_) {
	const int NUM_=400001;
	static ll fact[NUM_+1],factr[NUM_+1],inv[NUM_+1];
	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;
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>K;
	FOR(i,N) {
		A[i][0]=1;
		cin>>A[i][1];
		for(j=2;j<=300;j++) A[i][j]=A[i][j-1]*A[i][1]%mo;
		FOR(j,301) (S[j]+=A[i][j])%=mo;
	}
	
	for(i=1;i<=K;i++) {
		ll ret=0;
		FOR(j,i+1) (ret+=comb(i,j)*S[j]%mo*S[i-j])%=mo;
		ll tmp=S[i];
		FOR(j,i) tmp=2*tmp%mo;
		ret-=tmp;
		ret=(ret%mo+mo)*((mo+1)/2)%mo;
		cout<<ret<<endl;
	}
	
}

まとめ

まぁこれは典型かな…。