TopCoder SRM 540 Div1 Medium RandomColoring

問題

木の棒がN本、環状に並んでいる。
i番目とi+1番目、0番目とN-1番目はそれぞれ隣り合っている。


0番の木の棒を(startR, startG, startB)の色に塗る。
i番目の棒を塗り終えた後、i+1番目の棒を、次のようにして塗る。


i番目の棒の色を(R', G', B')、新しく塗る色を(R, G, B)とするとき、
0≦R<maxR
0≦G<maxG
0≦B<maxB
かつ、|R'-R|≦d2, |G'-G|≦d2, |B'-B|≦d2
かつ、|R'-R|, |G'-G|, |B'-B|のうちいずれかがd1以上である。


このような色のうちから等確率で一つが選ばれる。
N-1番目の棒まで色を塗ったとき、N-1番の棒と0番の棒が上の条件を満たしている確率はいくつか、求めよ。

制約条件

N≦40
maxR, maxG, maxB≦50

方針

(R, G, B)の状態からは、
(R-d2,R+d2)x(G-d2,G+d2)x(B-d2,B+d2)の立方体から、
(R-d1,R+d1)x(G-d1,G+d1)x(B-d1,B+d1)の立方体をくりぬいた状態のうち一つに等確率で遷移する。


この更新を愚直にやるとO(NRGBd2)の時間がかかってTLEするが、
足す区間の端にだけ値を足して、更新作業を一気にやることで、
足しこむ計算がO(1)ででき、間に合うようになる。


足しこむ計算は、一次元で説明すると次のようになる。

□□■■■■■■■□□
上の黒い範囲に等しくpを足したいとき、
□□+p■■■■■■-p□
とする。
その後で、左から順にa[i+1] += a[i]とすれば、
□□+p+p+p+p+p+p+p□□が残る。

この操作は最後に一度やるだけでいいので、直方体への足しこみがO(1)で出来ることになる。

ソースコード

int R, G, B;
double dp[2][64][64][64];

void add(double (*a)[64][64], int x, int X, int y, int Y, int z, int Z, double v){
	if(x > X || y > Y || z > Z) return;
	x = max(x, 0); y = max(y, 0); z = max(z, 0);
	X = min(X, R-1); Y = min(Y, G-1); Z = min(Z, B-1);
	X++; Y++; Z++;
	a[x][y][z] += v;
	a[X][y][z] -= v;
	a[x][Y][z] -= v;
	a[x][y][Z] -= v;
	a[X][Y][z] += v;
	a[X][y][Z] += v;
	a[x][Y][Z] += v;
	a[X][Y][Z] -= v;
}
int calc(int x, int X, int y, int Y, int z, int Z){
	if(x > X || y > Y || z > Z) return 0;
	x = max(x, 0); y = max(y, 0); z = max(z, 0);
	X = min(X, R-1); Y = min(Y, G-1); Z = min(Z, B-1);
	return (X - x + 1) * (Y - y + 1) * (Z - z + 1);
}
void fix(double (*a)[64][64]){
	rep(z,B+1) rep(y,G+1) rep(x,R) a[x+1][y][z] += a[x][y][z];
	rep(z,B+1) rep(x,R+1) rep(y,G) a[x][y+1][z] += a[x][y][z];
	rep(y,G+1) rep(x,R+1) rep(z,B) a[x][y][z+1] += a[x][y][z];
}

class RandomColoring {
	public:
	double getProbability(int N, int maxR, int maxG, int maxB, int startR,
 int startG, int startB, int d1, int d2) 
	{
		memset(dp, 0, sizeof(dp));
		R = maxR; G = maxG; B = maxB;
		double (*cur)[64][64] = dp[0], (*next)[64][64] = dp[1];
		
		cur[startR][startG][startB] = 1;
		rep(i,N-1){
			rep(j,R+1) rep(k,G+1) rep(l,B+1) next[j][k][l] = 0;
			rep(j,R+1) rep(k,G+1) rep(l,B+1) if(cur[j][k][l]){
				int x[4], y[4], z[4];
				x[0] = j-d2; x[1] = j-d1; x[2] = j+d1; x[3] = j+d2;
				y[0] = k-d2; y[1] = k-d1; y[2] = k+d1; y[3] = k+d2;
				z[0] = l-d2; z[1] = l-d1; z[2] = l+d1; z[3] = l+d2;
				
				int cnt = calc(x[0], x[3], y[0], y[3], z[0], z[3]) - 
					calc(x[1]+1, x[2]-1, y[1]+1, y[2]-1, z[1]+1, z[2]-1);
				if(!cnt) continue;
				double p = cur[j][k][l] / cnt;
				
				add(next, x[0], x[3], y[0], y[3], z[0], z[3], p);
				add(next, x[1]+1, x[2]-1, y[1]+1, y[2]-1, z[1]+1, z[2]-1, -p);
			}
			fix(next);
			swap(cur, next);
		}
		double ans = 0;
		rep(i,R+1) rep(j,G+1) rep(k,B+1){
			if(abs(i-startR) <= d2 && abs(j-startG) <= d2 && abs(k-startB) <= d2)
			if(abs(i-startR) >= d1 || abs(j-startG) >= d1 || abs(k-startB) >= d1)
			ans += cur[i][j][k];
		}
		return 1-ans;
	}
};