Codeforces #331 Div2 D. Wilbur and Trees

手抜きしてTLEした。
http://codeforces.com/contest/596/problem/D

問題

1列にN個の木が並んでいる。
木の高さはいずれもHで、各木の座標はX[i]とする。

この木の集合に風が吹くと、倒れていない左端または右端の木が確率1/2でいずれか倒れる(木が1本の時はそれが倒れる)。
また、倒れる方向は確率pで左側、(1-p)で右側となる。

木が倒れたとき、その方向でH未満の距離に別の木があると、その木も巻き添えを食って同じ方向に倒れる。
倒れた木は当然長さH分の土地を占める。ただし場合によっては木が倒れる位置が重なるので、H*Nの土地を占めるとは限らない。

全ての木が倒れたとき、木が占める土地の期待値を求めよ。

解法

木は両端からしか倒れないので、残された木の区間に対しメモ化再帰をすれば良い。
残された木に対し、以下の事象が生じえる。

• 左端の木が左に倒れる(確率p/2)
• 左端の木が右に倒れる(確率(1-p)/2)。この場合他の木を巻き添えにすることがある。
• 右端の木が左に倒れる(確率p/2)。この場合他の木を巻き添えにすることがある。
• 右端の木が右に倒れる(確率(1-p)/2)。

この際注意点がある。
例えば左端の木が左に倒れたとき、その左にすでに倒れた木があるとする。
その木が右に倒れていた場合、その木との距離が2H未満だと折り重なる可能性がる。
よってDPの際は、求める区間に左右隣接する木がどちらに倒れたかも状態として保持する必要がある。

また、木が1個倒れたときどこまでの木を巻き込むかは事前に計算しておこう。
DPの過程で毎回計算してしまうと、自分みたいにTLEする。

int N,H;
double P,Q;
int X[10101];
int le[10101],ri[10101];
double memo[2020][2020][2][2];

double dodo(int L,int R,int LL,int RL) {
	if(L>R) return 0;
	if(memo[L][R][LL][RL]>=0) return memo[L][R][LL][RL];
	double ret=0;
	
	int LX=X[L-1] + (LL==0)*H;
	int RX=X[R+1] - (RL==1)*H;
	
	if(L==R) {
		ret=P*min(H,X[L]-LX)+Q*min(H,RX-X[L]);
	}
	else {
		// left tree
		ret += P*(min(H,X[L]-LX)+dodo(L+1,R,1,RL));
		int NL=min(R,ri[L]);
		
		if(NL>=R) ret += Q*(min(RX,X[R]+H)-X[L]);
		else ret += Q*(X[NL]+H-X[L]+dodo(NL+1,R,0,RL));
		
		// right tree
		ret += Q*(min(H,RX-X[R])+dodo(L,R-1,LL,0));
		int NR=max(L,le[R]);
		if(NR<=L) ret += P*(X[R]-max(LX,X[L]-H));
		else ret += P*(X[R]-(X[NR]-H)+dodo(L,NR-1,LL,1));
		
		ret /= 2;
	}
	
	return memo[L][R][LL][RL]=ret;
}


void solve() {
	int i,j,k,l,r,x,y; string s;
	
	FOR(y,2010) FOR(x,2010) FOR(r,4) memo[y][x][r/2][r%2]=-1;
	
	cin>>N>>H>>P;
	Q=1-P;
	X[0]=-300000000;
	X[N+1]=300000000;
	FOR(i,N) cin>>X[i+1];
	sort(X+1,X+N+1);
	
	for(i=1;i<=N;i++) le[i]=(X[i]-X[i-1]<H)?le[i-1]:i;
	for(i=N;i>=1;i--) ri[i]=(X[i+1]-X[i]<H)?ri[i+1]:i;
	
	_P("%.12lf\n",dodo(1,N,0,1));
}

まとめ

シンプルで面白い問題でした。
でも手抜きTLEはいただけない…。