AtCoder ARC #028 : D - 注文の多い高橋商店

勉強になりました。
http://arc028.contest.atcoder.jp/tasks/arc028_4

問題

N種類の商品がそれぞれA[i]個あるので、合計でM個選びたい。

これだけでは単純なので、Q個のクエリを処理したい。
各クエリはK[i]、X[i]からなり、M個のうちK[i]の種類の商品はX[i]個ぴったり選ぶようにしたい。
各クエリの条件を満たす賞品の選び方を答えよ。

解法

まず前提として、このような重複組みあわせは蟻本にも記載があるように以下のように書ける。
i番目の商品まででj個選ぶ選びかたをdp[i+1][j]とすると
dp[i+1][j]=dp[i+1][j-1]-dp[i][j-1-A[i]]+dp[i][j]

部分点解法

上記DPを2回行っておく。
すなわち、商品を前から順に0～iまでをj個選ぶ選び方を上記式で求めておくと合わせ、逆向きに商品を(N-1)～kまでをl個選ぶ選び方を求めておく。

クエリK[i]、X[i]が来たら、0～K[i]-1番目までをp個、逆向きに(N-1)～K[i]+1番目までをM-X[i]-p個選べば、題意を満たす選び方ができる。
後はpを総当たりしていけばよい。
この方式だと、pの最大値がMなので1クエリあたりO(M)、全体でO(QM)かかる。

満点解法

上記式dp[i+1][j]=dp[i+1][j-1]-dp[i][j-1-A[i]]+dp[i][j]を変形するとdp[i][j]=dp[i+1][j]-dp[i+1][j-1]+dp[i][j-1-A[i]]である。

最初のDPは商品を前から順に行っていったが、実は商品の順はどうでもよい。
言い換えれば、上記変形式を用いるとN種類の商品を選ぶ選び方から1つ巻き戻すことができる。
すなわちdp2[i][j]=dp[N][j]-dp[N][j-1]+dp2[i][j-1-A[i]]と計算できる。

K[i]番の商品をX[i]個選ぶということは、それ以外の(N-1)個の商品をM-X[i]個選ぶということなので、dp2[K[i]][M-X[i]]が答えとなる。
先にdp2をO(NM)掛けて計算しておくと、各クエリはその値を読むだけなのでO(Q)ですむ。

int N,M,Q;
int A[3000];
ll mo=1000000007;
ll dp1[2005][2005],dp2[2005][2005];

void solve() {
	int i,j,k,l,r,x,y;
	
	cin>>N>>M>>Q;
	FOR(i,N) cin>>A[i];
	
	FOR(i,N+1) dp1[i][0]=1;
	FOR(i,N) FOR(j,M) {
		dp1[i+1][j+1]=dp1[i+1][j]+dp1[i][j+1];
		if(j-A[i]>=0) dp1[i+1][j+1]+=mo-dp1[i][j-A[i]];
		dp1[i+1][j+1]%=mo;
	}
	FOR(i,N) FOR(j,M+1) {
		dp2[i][j]=dp1[N][j];
		if(j) dp2[i][j]+=mo-dp1[N][j-1];
		if(j-1-A[i]>=0) dp2[i][j]+=dp2[i][j-1-A[i]];
		dp2[i][j]%=mo;
	}
	
	FOR(k,Q) {
		cin>>x>>y;
		_P("%lld\n",dp2[x-1][M-y]);
	}
}

まとめ

答えのコードは割とあっさり。この解法は勉強になった。