Submission #1143628

#	Time	Username	Problem	Language	Result	Execution time	Memory
1143628		Muhammet	Balanced Tree (info1cup18_balancedtree)	C++20	0 / 100	33 ms	2116 KiB

balancedtree

#include "bits/stdc++.h"

using namespace std;

#define SZ(s) (int)s.size()

const int N = 2e5 + 5;

int n, dp[5][N], a[N];

vector <vector <int>> v;

void f(int x, int y){
	vector <int> v1, v2;
	for(auto i : v[x]){
		if(i == y) continue;
		f(i, x);
		v1.push_back(min((!a[i] ? 0 : dp[1][i]) + 1, dp[1][x])), v2.push_back(min((a[i] ? 0 : dp[2][i]) + 1, dp[2][x]));
	}
	sort(v1.rbegin(), v1.rend()), sort(v2.rbegin(), v2.rend());
	if(SZ(v1)) dp[1][x] = v1.back(), v1.pop_back();
	if(SZ(v2)) dp[2][x] = v2.back(), v2.pop_back();
	if(SZ(v1)) dp[3][x] = v1.back(), v1.pop_back();
	if(SZ(v2)) dp[4][x] = v2.back(), v2.pop_back();

}

void f1(int x, int y){
	if(dp[1][y]-1 == dp[1][x]) dp[1][x] = min(dp[3][y] + 1, dp[1][x]);
	else dp[1][x] = min(dp[1][x], dp[1][y] + 1);
	if(dp[2][y]-1 == dp[2][x]) dp[2][x] = min(dp[4][y] + 1, dp[2][x]);
	else dp[2][x] = min(dp[2][x], dp[2][y] + 1);
	// dp[1][x] = min((!a[y] ? 0 : dp[1][y]) + 1, dp[1][x]), dp[2][x] = min((a[y] ? 0 : dp[2][y]) + 1, dp[2][x]);
	for(auto i : v[x]){
		if(i == y) continue;
		f1(i, x);
	}
}

void solve(){
	cin >> n;
	v.clear();
	v.resize(n+1);
	for(int i = 1; i < n; i++){
		int u1, u2;
		cin >> u1 >> u2;
		v[u1].push_back(u2);
		v[u2].push_back(u1);
	}

	dp[1][0] = dp[2][0] = dp[3][0] = dp[4][0] = 1e8;
	a[0] = 1e8;
	for(int i = 1; i <= n; i++){
		cin >> a[i];
		dp[1][i] = dp[2][i] = dp[3][i] = dp[4][i] = 1e8;
	}

	f(1, 0);
	f1(1, 0);

	int ans = 0, cnt0 = 0, cnt1 = 0;
	for(int i = 1; i <= n; i++){
		if(a[i]){
			cnt1++;
			ans = max(ans, dp[2][i]);
		}
		else {
			cnt0++;
			ans = max(ans, dp[1][i]);
		}
	}
	if(cnt1 == 1 or cnt0 == 1){
		cout << -1 << '\n';
		return;
	}
	cout << ans << '\n';
	for(int i = 1; i <= n; i++){
		cout << a[i] << ' ';
	}
	cout << '\n';
}

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

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

	return 0;
}

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