天下一プログラマーコンテスト2012 予選A D : アリの巣

問題概要

頂点数V(<2000)の最大次数3以下の根付き木が与えられる。根から蟻が出発し、各ノードで高々一つの特定の子があればそこへ降りる。行き止まりがあればその葉を潰して蟻は死ぬ。それ以外で(高々二つの)分岐があればどちらかに向かう。特定の葉までたどり着くのに必要な蟻の個数の期待値を求める問題。

解法

各ノードに対して、そこを通った蟻がk匹死んだときにゴールまでたどり着ける確率をdpで計算できる。更新式は各ノードで左右の部分木のサイズごとの積だけ計算する必要があるけど、このタイプの木DPは計算量がO(V^2)で抑えられているので十分間に合う。

acceptされたコード

#include <cstdio>
#include <cstdlib>
#include <algorithm>
#include <numeric>
#include <cmath>
#include <vector>
using namespace std;

const int MAX_W = 50;
const int MAX_N = MAX_W*MAX_W;
const int MAX_COMB = MAX_W*MAX_W;
double comb[MAX_COMB + 1][MAX_COMB + 1];

void prepareCombination() {
	for (int n = 0; n <= MAX_COMB; ++n) {
		comb[n][0] = comb[n][n] = pow(0.5, n);
		for (int k = 1; k < n; ++k) {
			comb[n][k] = (comb[n-1][k-1] + comb[n-1][k]) * 0.5;
		}
	}
}


const int dy[] = {1,-1,0,0}, dx[] = {0,0,1,-1};


int W, V, G;
char board[MAX_W][MAX_W + 1];
int nums[MAX_W][MAX_W];

bool visited[MAX_N][MAX_N];
double memo[MAX_N][MAX_N];
double dp[MAX_N][MAX_N];

vector<int> digraph[MAX_N];
vector<int> graph[MAX_N];

int size[MAX_N];
bool hasGoal[MAX_N];

void dfs(int y, int x, int dir) {
	visited[y][x] = true;
	int id = nums[y][x];
	if (id == G) {
		hasGoal[id] = true;
	}
	size[id] = 1;
	for (int k = 0; k < 4; ++k) {
		const int ny = y + dy[k], nx = x + dx[k];
		if ( 0 <= ny && ny < W && 0 <= nx && nx < W && board[ny][nx] == '.' && !visited[ny][nx]) {
			int nextId = nums[ny][nx];
			(dir == k ? digraph : graph)[id].push_back(nextId);
			dfs(ny, nx, k);
			if (hasGoal[nextId]) {
				hasGoal[id] = true;
			}
			size[id] += size[nextId];
		}
	}
}

void rec(int id) {
	//base
	if (id == G) {
		dp[id][0] = 1.0;
		return ;
	}
	if (!hasGoal[id]) {
		return ;
	}

	int add = 0;
	if (!digraph[id].empty()) {
		int tar = digraph[id][0];
		rec(tar);
		if (hasGoal[tar]) {
			for (int dead = 0; dead < V; ++dead) {
				dp[id][dead] = dp[tar][dead];
			}
			return ;
		}
		else {
			for (add = 0; add < size[tar]; ++add) {
				dp[id][add] = 0.0;
			}
		}
	}

	if ((int)graph[id].size() == 1) {
		int tar = graph[id][0];
		rec(tar);
		if (hasGoal[tar]) {
			for (int dead = 0; dead + add < add + size[tar]; ++dead) {
				dp[id][add+dead] = dp[tar][dead];
			}
		}
		// else 0.0
		return ;
	}

	{
		int t0 = graph[id][0], t1 = graph[id][1];
		rec(t0); rec(t1);
		if (hasGoal[t1]) {
			swap(t0, t1);
		}
		static double tp[MAX_N+1][MAX_N+1];
		for (int i = 0; i <= size[t0]; ++i) {
			for (int j = 0; j <= size[t1]; ++j) {
				tp[i][j] = 0.0;
			}
		}
		tp[0][0] = 1.0;
		for (int i = 0; i < size[t0]; ++i) {
			tp[i+1][0] = tp[i][0] * 0.5 * (1.0 - dp[t0][i]);
		}
		for (int j = 0; j < size[t1]; ++j) {
			tp[0][j+1] = tp[0][j] * 0.5;
		}
		for (int i = 0; i < size[t0]; ++i) {
			for (int j = 0; j < size[t1]; ++j) {
				tp[i+1][j+1] = (tp[i][j+1] * (j+1==size[t1] ? 1.0 : 0.5) * (1.0 - dp[t0][i]) +
						        tp[i+1][j] * (i+1==size[t0] ? 1.0 : 0.5));
			}
		}
		for (int i = 0; i <= size[t0]; ++i) {
			for (int j = 0; j <= size[t1]; ++j) {
				dp[id][add+i+j] += tp[i][j] * (j==size[t1] ? 1.0 : 0.5) * dp[t0][i];
			}
		}
		double cond = 1.0;
		for (int i = 0; i < add + size[t0] + size[t1]; ++i) {
			if (cond<1e-9) {
				dp[id][i] = 0.0;continue;
			}
			dp[id][i] /= cond;
			cond *= (1.0 - dp[id][i]);
		}
	}
}

double solve() {
	rec(0);
	double ret = 1.0;
	double notGoal = 1.0;
	for (int i = 0; i < V; ++i) {
		ret += notGoal * dp[0][i] * i;
		notGoal *= (1.0 - dp[0][i]);
	}
	return ret;
}

void init() {
	scanf("%d ", &W);
	if(W > 100){
		exit(-1);
	}
	for (int i = 0; i < W; ++i) {
		scanf(" %s ", board[i]);
	}
	for (int i = 0; i < W; ++i) {
		for (int j = 0; j < W; ++j) {
			if (board[i][j] == 'S' || board[i][j] == 'G' || board[i][j] == '.') {
				board[i][j] = '.';
				nums[i][j] = V++;
			}
		}
	}
	G = nums[W-1][W-1];
	dfs(0, 0, -1);
}

int main() {
	prepareCombination();
	init();
	printf("%.13f\n", solve());

	return 0;
}