TopCoder SRM 604 Div1 Medium FoxConnection

問題

木が与えられる。
それぞれの頂点には最大１匹のきつねがいる。
きつねは、最終的に以下の条件を満たすように移動したい。

• きつねがいる町が連結になっている
• 一つの町には最大一匹のきつねしかいない

このとき、きつねの移動距離の総和の最小値を求めよ。

制約条件

n≦50

方針

木DPするだけ。

きつねが木のノードrを頂点とする部分木にnum匹分布するためのコストをrec(r, num)とする。
で、その部分木以下に置くべききつねの数と、その部分木以下に最初にいるきつねの数には過不足があるかもしれないので、過不足は、全部rのノードにいるように調整する。

すなわち、子a以下にx匹を置くときにかかるコストは、
rec(a, x)に加えて、足りないときはrに待機してた余分なきつねがaへその分だけ補充しに、
多すぎるときはaで待機していたきつねがrにあふれてくると考えて、
rec(a, x) + abs(x - a以下のもともとのきつね)とすればDP全体で帳尻があう。

で、DPのを開始するルートを全通り試せばよい。

ソースコード

string fox;
int n, cnt[50], parent[50], memo[51][51];
vi e[50];

int initcost(int dist, int c, int p){
	int res = dist * (fox[c] == 'Y');
	rep(i, e[c].size()) if(e[c][i] != p) res += initcost(dist + 1, e[c][i], c);
	return res;
}
//c以下にnumの子を散らばらせるコストの最小値
int rec(int c, int p, int num){
	int &res = memo[c][num];
	if(res >= 0) return res;
	
	if(num == 0) return res = initcost(0, c, parent[c]);
	if(num == 1) return res = initcost(0, c, parent[c]);
	
	int dp[51][51], k = 0;
	rep(i, 51) rep(j, 51) dp[i][j] = inf;
	dp[0][0] = 0;
	
	rep(i, e[c].size()) if(e[c][i] != p){
		rep(j, 51) if(dp[k][j] < inf){
			rep(l, 51) if(j + l < 50) dp[k + 1][j + l] =
			min(dp[k + 1][j + l], dp[k][j] + rec(e[c][i], c, l) + abs(cnt[e[c][i]] - l));
		}
		k++;
	}
	return res = dp[k][num - 1];
}
int solve(int r){
	int cost = r == 0 ? 0 : initcost(1, parent[r], r);
	memset(memo, -1, sizeof(memo));
	return cost + rec(r, parent[r], cnt[0]);
}
void child(int c, int p){
	parent[c] = p;
	rep(i, e[c].size()) if(e[c][i] != p){
		child(e[c][i], c);
		cnt[c] += cnt[e[c][i]];
	}
}
class FoxConnection {
	public:
	int minimalDistance(vector <int> A, vector <int> B, string haveFox) {
		fox = haveFox;
		n = fox.size();
		rep(i, n){
			e[i].clear();
			cnt[i] = haveFox[i] == 'Y';
		}
		rep(i, n - 1){
			e[A[i] - 1].pb(B[i] - 1);
			e[B[i] - 1].pb(A[i] - 1);
		}
		int ans = inf;
		child(0, -1);
		rep(r, n) ans = min(ans, solve(r));
		return ans;
	}
};