Submission #1143567

#	Time	Username	Problem	Language	Result	Execution time	Memory
1143567		stdfloat	Balanced Tree (info1cup18_balancedtree)	C++20	13 / 100	48 ms	6592 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);

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

		if (!~a.ss.ff || a.ss.ss > (!c[i] ? 0 : dp[i][0]) + 1) a.ss = {i, (!c[i] ? 0 : dp[i][0]) + 1};
		if (!~a.ff.ff || a.ff.ss > a.ss.ss) swap(a.ff, a.ss);

		if (!~b.ss.ff || b.ss.ss > (c[i] ? 0 : dp[i][1] + 1)) b.ss = {i, (c[i] ? 0 : dp[i][1]) + 1};
		if (!~b.ff.ff || b.ff.ss > b.ss.ss) swap(b.ff, b.ss);
	}

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

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

		int ndp0 = (a.ff.ff != i ? a.ff.ss : a.ss.ss), ndp1 = (b.ff.ff != i ? b.ff.ss : b.ss.ss);
		dfs2(i, min(min(dp0, ndp0) + 1, (int)1e8), min(min(dp1, ndp1) + 1, (int)1e8), 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;

	vector<int> remc = c;

	vector<int> ans;
	int mn = INT_MAX;
	for (int z = 0; z < 2; z++) {
		c = remc;
		for (auto &i : c)
			if (!~i) i = z;

		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 (mn > mx) {
			mn = mx;
			ans = c;
		}
	}

 	if (mn == (int)1e8) cout << "-1\n";
 	else {
 		cout << mn << '\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: 13 / 13
• Subtask 3: 0 / 26
• Subtask 4: 0 / 45
• Subtask 5: 0 / 6