AtCoder ARC #104 : D - Multiset Mean

力技で解いてしまった。
https://atcoder.jp/contests/arc104/tasks/arc104_d

問題

1～Nの整数を最大K個ずつ選択するとき、その平均がXとなるのは何通りか。
X=1～Nについてそれぞれ答えよ。

解法

O(N^4*K)のゴリ押し解なのでご参考。
Xを総当たりするが、対称性から平均がXと(N+1)-Xの場合は組み合わせが同じなので、Xは半分だけ求める。
Xを定めたら、(1-X)～(N-X)をK個ずつ選ぶ場合の総和に対し、組み合わせ数をDPで求めて行こう。
総和の範囲を適切に絞ると計算回数が減って何とか間に合う。

int N,K;
ll mo;

ll from[5050*100+100];
ll to[5050*100+100];
ll ret[101];
void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>K>>mo;
	
	
	for(x=1;x<=(N+1)/2;x++) {
		ZERO(from);
		ZERO(to);
		int mi=x*(x-1)/2*K;
		int ma=(N-x)*(N-x+1)/2*K;
		int cur=mi+50;
		from[cur]=K+1;
		mi=ma=cur;
		for(i=x-1;i>=1;i--) {
			int pma=ma;
			
			mi=min(mi,mi+K*(i-x));
			ma=max(ma,ma+K*(i-x));
			ll sum[101]={};
			for(y=ma;y>=mi;y--) {
				to[y]=from[y];
				if(y+(x-i)<=pma) to[y]+=to[y+(x-i)];
				if(y-(K+1)*(i-x)<=pma) to[y]+=mo-from[y-(K+1)*(i-x)];
				while(to[y]>=mo) to[y]-=mo;
			}
			swap(from,to);
		}
		for(i=x+1;i<=N;i++) {
			int pmi=mi;
			for(y=mi;y<=cur;y++) {
				to[y]=from[y];
				if(y+(x-i)>=mi) to[y]+=to[y+(x-i)];
				if(y-(K+1)*(i-x)>=pmi) to[y]+=mo-from[y-(K+1)*(i-x)];
				while(to[y]>=mo) to[y]-=mo;
			}
			swap(from,to);
		}
		ret[x]=(from[cur]+mo-1)%mo;
		ret[(N+1)-x]=ret[x];
	}
	
	for(i=1;i<=N;i++) cout<<ret[i]<<endl;
	
}

まとめ

手元で3.6sかかったけどsubmitしたら2.6s程度で時間的にはよかった。
ただ微妙な配列処理のミスで2WAしたのは頂けない。