AtCoder ARC #119 : F - AtCoder Express 3

なるほど。
https://atcoder.jp/contests/arc119/tasks/arc119_f

問題

0～N番の(N+1)個の駅が並んでいる。
これらの間は3種類の列車が運行されている。
普通列車は、すべての駅で停車する。
光速列車と準急列車は、0番とN番の駅はいずれも停車するが、それ以外の駅はいずれか片方のみが停車する。

各列車は停車駅の間を時間1毎に運行し、時間1で移動できる。
1～(N-1)番の駅において、光速列車と準急列車のどちらが停車するか、もしくはどちらも指定されないかが入力で指示される。
指定されない駅において、光速列車と準急列車のどちらが停車するか総当たりしたとき、0番からN番の駅に移動する最短時間の総和を求めよ。

解法

各駅の停車する列車が確定したときにかかる最短時間を考える。
駅を戻らないケースを考えると、直前の駅から各駅で移動するか、同じ種別の列車の直前の駅から移動するかを考えればよく、それほど複雑な遷移にはならない。
ただし、実際には途中駅を1つ戻る方が良いケースがある。
そこで以下のように考える。

f(p,a,b,v) := p番目の駅に至るとき、光速列車での最寄駅までの最短時間がa、準急は同様にbで、直前に乗った列車がv(光速か準急の2値)であるような組み合わせの数
として、状態遷移を考える。
p+1番目の駅が光速の場合、

• f(p,a,b,0) → f(p+1,a+1,b,0) (直前も光速なら、そのまま光速でも普通でも進むのが最速)
• f(p,a,b,1) → f(p+1,min(a,b)+1,min(b,a+2)+1,0) (愚直に普通列車で進むか、直前の光速列車のから進むか。また、最寄の準急列車の駅までの時間も更新される)

こう考えるとf(N,*,*,*)はO(N^3)で求められるが、これだとTLEする。
重要なのは上記minの部分で、aとbが3以上離れるケースは生じないことを用いると、状態空間がO(N^2)に抑えられる。

int N,K;
string S;
const ll mo=1000000007;
int dp[4040][4040][5][2];

void go(int i,int nx,int ny,int nz,int x,int y,int z) {
	if(nx>ny+2) nx=ny+2;
	if(ny>nx+2) ny=nx+2;
	(dp[i+1][nx][ny-nx+2][nz]+=dp[i][x][y-x+2][z])%=mo;
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N>>K>>S;
	dp[0][0][2][0]=1;
	FOR(i,N) {
		FOR(x,N) for(y=max(0,x-2);y<=x+2;y++) {
			if(S[i]=='A'||S[i]=='?') {
				go(i,x+1,y,0,x,y,0);
				go(i,min(x,y)+1,min(x+2,y),0,x,y,1);
			}
			if(S[i]=='B'||S[i]=='?') {
				go(i,x,y+1,1,x,y,1);
				go(i,min(x,y+2),min(x,y)+1,1,x,y,0);
			}
		}
	}
	
	ll ret=0;
	FOR(x,N) FOR(y,5) if(x<K||(x+y-2)<K) ret+=dp[N-1][x][y][0]+dp[N-1][x][y][1];
	cout<<ret%mo<<endl;
}

まとめ

覚えておくべきテクだな。