Submission #497044

#	Time	Username	Problem	Language	Result	Execution time	Memory
497044		abc864197532	Balanced Tree (info1cup18_balancedtree)	C++17	0 / 100	615 ms	524292 KiB

balancedtree

#include <bits/stdc++.h>
using namespace std;
#define lli long long int
#define mp make_pair
#define eb emplace_back
#define pb push_back
#define X first
#define Y second
#define pii pair<int,int>
#define all(x) x.begin(), x.end()
#define rall(x) x.rbegin(), x.rend()
void abc() {cout << endl;}
template <typename T, typename ...U> void abc(T i, U ...j) {
	cout << i << ' ', abc(j...);
}
template <typename T> void printv(T l, T r) {
	for (; l != r; ++l) 
		cout << *l << " \n"[l + 1 == r];
}
#ifdef Doludu
#define test(x...) abc("[" + string(#x) + "]", x);
#else
#define test(x...) void(0);
#endif
const int N = 100000;

int main () {
	ios::sync_with_stdio(false);
	cin.tie(0);
	int t;
	cin >> t;
	while (t--) {
		int n;
		cin >> n;
		vector <vector <int>> adj(n);
		for (int i = 0, u, v; i < n - 1; ++i) {
			cin >> u >> v, --u, --v;
			adj[u].pb(v), adj[v].pb(u);
		}
		vector <int> a(n);
		for (int i = 0; i < n; ++i)
			cin >> a[i], assert(a[i] != -1);
		if (count(all(a), 0) == 1 || count(all(a), 1) == 1) {
			cout << -1 << endl;
			continue;
		}
		vector <vector <int>> dp1(n, vector <int>(2, n + 1));
		vector <vector <int>> dp2(n, vector <int>(2, n + 1));
		auto upd = [&](int i, int j, int v) {
			dp1[i][j] = min(dp1[i][j], v);
		};
		function<void(int, int)> dfs = [&](int v, int pa) {
			for (int u : adj[v]) if (u != pa) {
				dfs(u, v);
				upd(v, 0, dp1[u][0] + 1);
				upd(v, 1, dp1[u][1] + 1);
				upd(v, a[u], 1);
			}
		};
		function<void(int, int)> dfs2 = [&](int v, int pa) {
			if (~pa) {
				dp2[v][t] = dp2[pa][t] + 1;
				for (int t : {0, 1}) {
					dp2[v][t] = min(dp2[v][t], dp1[pa][t] + 1);
				}
				dp2[v][a[pa]] = min(dp2[v][a[pa]], 1);
			}
			for (int u : adj[v]) if (u != pa) {
				dfs2(u, v);
			}
		};
		dfs(0, -1);
		dfs2(0, -1);
		int ans = 0;
		for (int i = 0; i < n; ++i) {
			ans = max(ans, min(dp1[i][a[i]], dp2[i][a[i]]));
		}
		cout << ans << endl;
		printv(all(a));
	}
}

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