Submission #1143544

#	Time	Username	Problem	Language	Result	Execution time	Memory
1143544		stdfloat	Balanced Tree (info1cup18_balancedtree)	C++20	0 / 100	30 ms	2884 KiB

balancedtree

#include <bits/stdc++.h>
using namespace std;

using ll = long long;

#define ff	first
#define ss	second
#define pii	pair<int, int>

vector<int> c;

vector<vector<int>> E, dp;

void dfs1(int x, int p = -1) {
	for (auto i : E[x]) {
		if (i == p) continue;
		
		dfs1(i, x);

		dp[x][0] = min(dp[x][0], (!c[i] ? 0 : dp[i][0]) + 1);
		dp[x][1] = min(dp[x][1], (c[i] ? 0 : dp[i][1]) + 1);
	}
}

void dfs2(int x, int dp0 = (int)1e8, int dp1 = (int)1e8, int p = -1) {
	dp[x][0] = min(dp[x][0], dp0);
	dp[x][1] = min(dp[x][1], dp1);

	pii a = {-1, -1}, b = a;
	for (auto i : E[x]) {
		if (i == p) continue;

		if (!~a.ss || dp[a.ss][0] > dp[i][0]) a.ss = i;
		if (!~a.ff || dp[a.ff][0] > dp[a.ss][0]) swap(a.ff, a.ss);

		if (!~b.ss || dp[b.ss][1] > dp[i][1]) b.ss = i;
		if (!~b.ff || dp[b.ff][1] > dp[b.ss][1]) swap(b.ff, b.ss);
	}

	(c[x] ? dp1 : dp0) = 0;

	for (auto i : E[x]) {
		if (i == p) continue;

		int ndp0 = (a.ff != i ? a.ff : a.ss), ndp1 = (b.ff != i ? b.ff : b.ss);
		dfs2(i, min(dp0, (~ndp0 ? dp[ndp0][0] : (int)1e8)) + 1, min(dp1, (~ndp1 ? dp[ndp1][0] : (int)1e8)) + 1, x);
	}
}

void solve() {
	int n;
	cin >> n;

	E.assign(n, {});
	for (int i = 1; i < n; i++) {
		int x, y;
		cin >> x >> y; x--; y--;

		E[x].push_back(y);
		E[y].push_back(x);
	}

	c.assign(n, 0);
	for (auto &i : c)
		cin >> i;

	dp.assign(n, vector<int>(2, (int)1e8));
 	dfs1(0);

 	dfs2(0);

 	int mx = 0;
 	for (int i = 0; i < n; i++)
 		mx = max(mx, dp[i][c[i]]);

 	if (mx == (int)1e8) cout << "-1\n";
 	else {
 		cout << mx << '\n';
 		for (auto i : c)
 			cout << i << ' ';
 		cout << '\n';
 	}
}

int main() {
	ios::sync_with_stdio(false); cin.tie(nullptr);

	int T;
	cin >> T;
	while (T--) solve();
}

• Subtask 1: 0 / 10
• Subtask 2: 0 / 13
• Subtask 3: 0 / 26
• Subtask 4: 0 / 45
• Subtask 5: 0 / 6