kmjp's blog

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

Codeforces #225 Div1. D. Antimatter

テクニックを知っている人には簡単な問題。
http://codeforces.com/contest/383/problem/D

問題

N種類のデバイスがあり、それぞれA[i]個だけ物質または反物質を生成することができる。
これらのデバイスから1つ以上の連続したsubarrayを選択し、それぞれのデバイスの生成物を物質か反物質かを指定して、物質と反物質の合計が等しくなるようなsubarray及び生成物の選び方を答えよ。

解法

Editorialはずいぶんややこしいことが書いてあるが、実際は割とオーソドックスに回答することができる。
自力では解けなかったが、他の人の解答を参考に解いてみた。

各デバイスの生成物の総和をSとする。
まず、subarrayが先に定まっている場合、単純なDPで物質と反物質の差分数を管理すると、O(NS)でDPが完了する。
N<=1000、S<=10000よりこれは問題なく終わる。

問題は、subarrayの範囲も変化することである。
subarrayの数だけDPを行うと、subarrayの選び方がO(N^2)あるので全体でO(N^3S)かかって破たんする。
なんとか計算量を抑えるために、以下のようにすればよい。

  • 1番目のデバイスからDPを行うが、通常0番のデバイスを仮定してその選び方を1通りとして始める。しかし、各i+1番目のデバイスをDPする際、その前にi番目のデバイスまで選ばない、という選び方を1通り加えてやればよい。
    • これにより、DPを開始する位置を2つ目以降のデバイスにするときの数も考慮できる。
  • subarrayの終わりも同様で、i番目のデバイスまでDPした段階で、その時物質と反物質の差分が0の組み合わせをその都度答えに加えればよい。

合わせると以下のようにあっさり書ける。

ll mo=1000000007;
ll dpdp[1001][20001];
int N;

void solve() {
	int f,i,j,k,l,x,y,a;
	ll ret=0;
	
	cin>>N;
	FOR(i,N) {
		cin>>a;
		dpdp[i][10000]++;
		for(x=0;x<=20000;x++) {
			if(x+a<=20000) dpdp[i+1][x+a]=(dpdp[i+1][x+a]+dpdp[i][x])%mo;
			if(x-a>=0) dpdp[i+1][x-a]=(dpdp[i+1][x-a]+dpdp[i][x])%mo;
		}
		ret += dpdp[i+1][10000];
	}
	cout << ret%mo << endl;
}

まとめ

各subarrayに対する組み合わせ数の総計を求めるとき、この毎回組み合わせ数に1を足すテクは応用できそうだ。