kmjp's blog

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

CSAcademy Round #32 : G. Sum of Powers

これは勉強になった。
https://csacademy.com/contest/round-32/task/sum-of-powers/

問題

正整数N,M,Kが与えられる。
Nを正整数K個に分割したmultiset aを考える。すべてのaにおいて \sum_i {a_i}^Mの総和を答えよ。

解法

以下を考える。
f(x,y) := a中で整数xがちょうどy個あるような組み合わせの数

すると求める解は \displaystyle \sum_x \sum_y \left( f(x,y)\times y \times x^M \right)となる。
あとはf(x,y)を求めることを考えよう。
直接f(x,y)を求めるのは難しいので、以下を求める。
g(x,y) := a中で整数xがy個以上あるような組み合わせの数
すると f(x,y) = g(x,y) - g(x,y+1)で計算できる。

次にg(x,y)の求め方だが、残る(K-y)要素で合計(N-x*y)の値を作ればよいので、これは分割数のDPで求められる。

int N,M,K;
ll mo=1000000007;

ll sp[4098][4098];

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	sp[0][0]=sp[1][1]=1;
	for(x=1;x<=4096;x++) {
		for(y=1;y<=4096;y++) if(sp[x][y]) {
			if(y+x<=4096) {
				sp[x][y+x]+=sp[x][y];
				if(sp[x][y+x]>=mo) sp[x][y+x]-=mo;
			}
			if(y+1<=4096) {
				sp[x+1][y+1]+=sp[x][y];
				if(sp[x+1][y+1]>=mo) sp[x+1][y+1]-=mo;
			}
		}
	}
	
	cin>>N>>K>>M;
	
	ll ret=0;
	for(i=1;i<=N;i++) {
		ll v=1;
		FOR(j,M) v=v*i%mo;
		for(j=1;j<=K && i*j<=N;j++) {
			ll pat = sp[K-j][N-i*j];
			if(j+1<=K && i*(j+1)<=N) pat += mo- sp[K-(j+1)][N-i*(j+1)];
			(ret += v*j%mo*pat)%=mo;
		}
	}
	cout<<ret<<endl;
}

まとめ

包除原理は思いつかなかった。