SRM 540 550pt : RandomColoring

問題概要

RGBの値が、直前の値と離れすぎず近づきすぎずになるようにランダムに塗っていく。最初と最後が整合しない確率を求める問題(適当)。

解法

愚直にやると、状態数が40*50^3で参照先が50^3となりTLEする。累積和を用いてO(1)で参照できるように加速すればよい。

acceptされたコード

#include <cstring>
#include <algorithm>
using namespace std;

const int MAX_N = 40;
const int MAX_C = 50;

int lowerR[MAX_C][2], lowerG[MAX_C][2], lowerB[MAX_C][2];
int upperR[MAX_C][2], upperG[MAX_C][2], upperB[MAX_C][2];

double dp[MAX_N + 1][MAX_C][MAX_C][MAX_C];
double accum[MAX_C + 1][MAX_C + 1][MAX_C + 1];

//[x1, x2)*[y1, y2)*[z1, z2)の範囲の和を返す
double get_range_sum(int x1, int x2, int y1, int y2, int z1, int z2){
	return accum[x2][y2][z2] - ( accum[x1][y2][z2] + accum[x2][y1][z2] + accum[x2][y2][z1] )
			+ ( accum[x2][y1][z1] + accum[x1][y2][z1] + accum[x1][y1][z2] ) - accum[x1][y1][z1];
}

struct RandomColoring {

	double getProbability(int N, int maxR, int maxG, int maxB, int startR, int startG, int startB, int d1, int d2) {

		//前処理
		for(int r=0; r<maxR; ++r){
			if(d1 > 0){
				upperR[r][0] = min(r + d1, maxR);
				lowerR[r][0] = max(r - d1 + 1, 0);
			}
			upperR[r][1] = min(r + d2 + 1, maxR);
			lowerR[r][1] = max(r - d2, 0);
		}

		for(int g=0; g<maxG; ++g){
			if(d1 > 0){
				upperG[g][0] = min(g + d1, maxG);
				lowerG[g][0] = max(g - d1 + 1, 0);
			}
			upperG[g][1] = min(g + d2 + 1, maxG);
			lowerG[g][1] = max(g - d2, 0);
		}

		for(int b=0; b<maxB; ++b){
			if(d1 > 0){
				upperB[b][0] = min(b + d1, maxB);
				lowerB[b][0] = max(b - d1 + 1, 0);
			}
			upperB[b][1] = min(b + d2 + 1, maxB);
			lowerB[b][1] = max(b - d2, 0);
		}

		//DP
		dp[0][startR][startG][startB] = 1.0;
		for(int i=0; i<N - 1; ++i){
			//累積和を計算する。
			for(int r=0; r<maxR; ++r){
				for(int g=0; g<maxG; ++g){
					for(int b=0; b<maxB; ++b){
						int total = (upperR[r][1] - lowerR[r][1]) * (upperG[g][1] - lowerG[g][1]) * (upperB[b][1] - lowerB[b][1])
								- (upperR[r][0] - lowerR[r][0]) * (upperG[g][0] - lowerG[g][0]) * (upperB[b][0] - lowerB[b][0]);
						double p = total > 0 ? 1.0 / total : 0.0;
						accum[r+1][g+1][b+1] = (accum[r][g+1][b+1] + accum[r+1][g][b+1] + accum[r+1][g+1][b])
								- (accum[r+1][g][b] + accum[r][g+1][b] + accum[r][g][b+1]) + accum[r][g][b];
						accum[r+1][g+1][b+1] += dp[i][r][g][b] * p;
					}
				}
			}

			//dpテーブルを更新する
			for(int r=0; r<maxR; ++r){
				for(int g=0; g<maxG; ++g){
					for(int b=0; b<maxB; ++b){
						dp[i+1][r][g][b] += get_range_sum(lowerR[r][1], upperR[r][1], lowerG[g][1], upperG[g][1], lowerB[b][1], upperB[b][1]);
						if(d1 > 0){
							dp[i+1][r][g][b] -= get_range_sum(lowerR[r][0], upperR[r][0], lowerG[g][0], upperG[g][0], lowerB[b][0], upperB[b][0]);
						}
						if(dp[i+1][r][g][b] < 1e-100){
							dp[i+1][r][g][b] = 0.0;
						}
					}
				}
			}
		}

		double ans = 0.0;
		for(int r=0; r<maxR; ++r){
			for(int g=0; g<maxG; ++g){
				for(int b=0; b<maxB; ++b){
					if(abs(startR - r) <= d2 && abs(startG - g) <= d2 && abs(startB - b) <= d2 &&
							(abs(startR - r) >= d1 || abs(startG - g) >= d1 || abs(startB - b) >= d1) ){
						ans += dp[N-1][r][g][b];
					}
				}
			}
		}

		return 1.0 - ans;
	}

};