Codeforces #627 div3 F - Maximum white subtree (graph dp, reroot)

https://codeforces.com/contest/1324/problem/F

reroot 첫 성공

 

<문제설명>

무방향트리가 주어진다. 각 정점마다 하얀색(1)또는 검은색(0)으로 색칠되어 있다.

i번 정점을 포함하는 서브트리 중 하얀색정점의 개수 - 검은색 정점의 개수가 최대인 경우 그 값을 각각 구해라 (i=1,2,3,4...n)

 

<풀이>

1번 정점을 root로 하는 전체 트리에서

dp[u] : u를 포함한 서브트리에서 cntw-cntb 최대값

     

u하양인경우 1 검정인경우 -1

 

u에 인접한 v에 대해 ans[v] = dp[v] + max(dp[u],0)이다.

v에서원하는 값 ans[v]가 u,v에대한 상태로만 정의되므로 reroot 적용가능하다.

 

dp[u] -= max(dp[v],0)

dp[v] += max(dp[u],0)

go(u,v)

dp[u] = memo1;

dp[v] = memo2;

 

백트래킹해주면된다.

 

<코드>

int n, dp[MAX], ans[MAX];
bool white[MAX];
vector<int> adj[MAX];
int makeDp(int u, int p) {
	dp[u] = white[u] ? 1 : -1;
	for (int v : adj[u]) if (v != p) 
		dp[u] += max(0, makeDp(v, u));
	return dp[u];
}
void reroot(int u, int p) {
	ans[u] = dp[u];

	for (int v : adj[u]) if (v != p) {
		int t1 = dp[u], t2 = dp[v];
		dp[u] -= max(dp[v], 0);
		dp[v] += max(dp[u], 0);
		reroot(v, u);
		dp[u] = t1;
		dp[v] = t2;
	}
}

int main() {
	FAST; cin >> n;
	for (int i = 1; i <= n; i++) cin >> white[i];
	for (int i = 0, u, v; i < n - 1; i++) {
		cin >> u >> v;
		adj[u].push_back(v);
		adj[v].push_back(u);
	}
	makeDp(1, 0);
	reroot(1, 0);
	for (int i = 1; i <= n; i++) cout << ans[i] << " \n"[i==n];
}