Codeforces #408 Div2 E. Exam Cheating

もうひと押しだった。
http://codeforces.com/contest/796/problem/E

問題

N問の問題を順に解くことを考える。
両隣には別の生徒がおり、それぞれが正解する問題の一覧が与えられる。

両者の答えをカンニングし、自身の正答数を最大化したい。
自身は両者を合計P回カンニングできる。1回のカンニングでは片方の生徒に対し最大でK個の連続する問題を見ることができる。
同時に両隣の生徒のカンニングをしてもよい。

自身の正答数の最大値を求めよ。

解法

以下を考える。
f(i,a,b,x) := i問目までの問題で、直前にカンニングした両者の問題番号はa,bであり、残りx回カンニング可能な状態で得られる最大正答数

1人目の生徒の(a+1)～(i+1)番目および2人目の生徒の(b+1)～(i+1)番目をカンニングするかどうかを考える。
するとf(i,a,b,x)からf(i+1,a,b,x),f(i+1,i+1,b,x-1),f(i+1,a,i+1,x-1),f(i+1,i+1,i+1,x-2)に遷移することがわかる。
(両者のカンニング範囲に対する追加正答数は先に前処理で計算しておくとよい)

最大で1回のカンニングでK問までしか回答がわからないため、i-aおよびi-bはK以下のケースだけ考えればよい。
ただ、それでもfの引数の組み合わせはO(N*P*K^2)かかりこのままではTLEする。
ただ、もともとP*K≧2Nであれば全問カンニング可能であるため、P*K＜2Nのケースのみ上記DPを行えばよい。
その場合計算量はO(NPK)で何とか間に合う。
ただ、Pが大きい場合とKが大きい場合、いずれもMLEしないためにはメモリレイアウトの工夫が必要。

以下のコードは、負のindexの扱いが面倒なため、後ろの問題から順に前に巻き戻す方向でDPしている。

int N,P,K;
int A[2][1110];
int num[2][1110][1110];
int only[2][1110][1110];
const int sz=500000;
int from[sz],to[sz];

int id(int p,int x,int y) {return p*(K+2)*(K+2)+x*(K+2)+y;}
void update(int p,int x,int y,int v) {
	int& r=to[id(p,x,y)];
	r = max(r,v);
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>P>>K;
	cin>>x;
	FOR(i,x) cin>>r, A[0][r-1]=1;
	cin>>x;
	FOR(i,x) cin>>r, A[1][r-1]=1;
	
	if(P/2*K>=N) {
		int cnt=0;
		FOR(i,N) cnt += A[0][i] | A[1][i];
		return _P("%d\n",cnt);
	}
	
	FOR(x,N) {
		for(y=x;y<N+K;y++) {
			num[0][x][y] = (y?num[0][x][y-1]:0) + A[0][y];
			num[1][x][y] = (y?num[1][x][y-1]:0) + A[1][y];
			only[0][x][y] = (y?only[0][x][y-1]:0) + (A[0][y] & !A[1][y]);
			only[1][x][y] = (y?only[1][x][y-1]:0) + (A[1][y] & !A[0][y]);
		}
	}
	
	for(i=N-1;i>=0;i--) {
		ZERO(to);
		for(j=0;j<=P;j++) {
			FOR(x,K+1) FOR(y,K+1) {
				int r = from[id(j,x,y)];
				update(j,min(x+1,K),min(y+1,K),r);
				if(j) {
					if(x<=y) {
						update(j-1,1,min(y+1,K),r+num[0][i][i+x-1]);
						update(j-1,min(x+1,K),1,r+num[1][i][i+x-1]+only[1][i+x][i+y-1]);
					}
					else {
						update(j-1,1,min(y+1,K),r+num[0][i][i+y-1]+only[0][i+y][i+x-1]);
						update(j-1,min(x+1,K),1,r+num[1][i][i+y-1]);
					}
					
				}
				if(j>=2) {
					if(x<=y) update(j-2,1,1,r+num[0][i][i+x-1]+only[1][i][i+y-1]);
					else update(j-2,1,1,r+only[0][i][i+x-1]+num[1][i][i+y-1]);
				}
			}
		}
		
		swap(from,to);
	}

	int ma=0;
	FOR(i,sz) ma=max(ma,from[i]);
	cout<<ma<<endl;
	
}

まとめ

うーん、P*K≦2Kを利用するのは思いついたけど、その後のDPで混乱してしまった。