Codeforces_forever Round #1 (Div3) F

//решение
Эта задача на технику «переподвешивания». Сначала давайте посчитаем ответ для какого-то фиксированного корня. Как мы можем это сделать? Пусть dpv равно максимально возможной разницы между количеством белых и черных вершин в каком-то таком поддереве v (да, поддерево коревого дерева, то есть v и все ее прямые и косвенные дети), что содержит вершину v. Мы можем посчитать это динамическое программирование при помощи простого dfs, для вершины v оно будет выглядеть следующим образом: dpv=av+∑to∈children(v)max(0,dpto).

Окей, мы где-то сохранили ответ для корня . Что дальше? Давайте попробуем сменить корень с вершины v на какую-то смежную с ней вершину to. Какие состояния динамического программирования изменятся? Только dpv и dpto. Сначала нам необходимо «удалить» ребенка to из поддерева вершины v: dpv=dpv−max(0,dpto). Затем нам необходимо «прицепить» вершину v и сделать ее сыном вершины to: dpto=dpto+max(0,dpv). Затем нам нужно запустить этот процесс рекурсивно из вершины to (сохранить ответ, переподвесить дерево и так далее), и когда он закончится, нам нужно будет «откатить» наши изменения. Теперь v опять является корнем и мы можем попробовать следующего сына v в качестве корня.

Асимптотика решения: O(n).
//реализация
#include <bits/stdc++.h>

using namespace std;

vector<int> a;
vector<int> dp;
vector<int> ans;
vector<vector<int>> g;

void dfs(int v, int p = -1) {
	dp[v] = a[v];
	for (auto to : g[v]) {
		if (to == p) continue;
		dfs(to, v);
		dp[v] += max(dp[to], 0);
	}
}

void dfs2(int v, int p = -1) {
	ans[v] = dp[v];
	for (auto to : g[v]) {
		if (to == p) continue;
		dp[v] -= max(0, dp[to]);
		dp[to] += max(0, dp[v]);
		dfs2(to, v);
		dp[to] -= max(0, dp[v]);
		dp[v] += max(0, dp[to]);
	}
}

int main() {
#ifdef _DEBUG
	freopen("input.txt", "r", stdin);
//	freopen("output.txt", "w", stdout);
#endif

	int n;
	cin >> n;
	a = dp = ans = vector<int>(n);
	g = vector<vector<int>>(n);
	for (int i = 0; i < n; ++i) {
		cin >> a[i];
		if (a[i] == 0) a[i] = -1;
	}
	for (int i = 0; i < n - 1; ++i) {
		int x, y;
		cin >> x >> y;
		--x, --y;
		g[x].push_back(y);
		g[y].push_back(x);
	}

	dfs(0);
	dfs2(0);
	for (auto it : ans) cout << it << " ";
	cout << endl;

	return 0;
}