#include<bits/stdc++.h>
#define pb emplace_back
#define AI(i) begin(i), end(i)
using namespace std;
using ll = long long;
#ifdef KEV
#define DE(args...) kout("[ " + string(#args) + " ] : ", args)
void kout() { cerr << endl; }
template<class T, class ...U>
void kout(T a, U ...b) { cerr << a << ' ', kout(b...); }
void debug(auto L, auto R) { while (L != R) cerr << *L << " \n"[next(L) == R], ++L; }
#else
#define debug(...) 0
#define DE(...) 0
#endif
const int MAX_N = 500010, inf = 1e9;
int n, c[MAX_N];
vector<int> edge[MAX_N];
void clear() {
for (int i = 1; edge[i].size() ; ++i)
edge[i].clear();
}
int dis[MAX_N], sz[MAX_N];
bool ban[MAX_N];
int dfs_sz(int x, int lst = -1) {
sz[x] = 1;
for (int u : edge[x]) if (u != lst && !ban[u])
sz[x] += dfs_sz(u, x);
return sz[x];
}
int dfs_cen(int x, int allsz, int lst = -1) {
for (int u : edge[x]) if (u != lst && !ban[u])
if (sz[u] * 2 > allsz)
return dfs_cen(u, allsz, x);
return x;
}
void chmin(int &v, int val) { v = min(v, val); }
int min_dep[2];
void dfs_update(int x, int lst = -1) {
static int d;
if (ban[x]) return;
++d;
chmin(dis[x], min_dep[c[x]] + d);
chmin(min_dep[c[x]], d);
for (int u : edge[x]) if (u != lst)
dfs_update(u, x);
--d;
}
void dfs_solve(int side) {
if (ban[side]) return;
int cen = dfs_cen(side, dfs_sz(side));
min_dep[0] = min_dep[1] = inf;
min_dep[ c[cen] ] = 0;
ban[cen] = true;
for (int u : edge[cen])
dfs_update(u);
min_dep[0] = min_dep[1] = inf;
reverse(AI(edge[cen]));
for (int u : edge[cen])
dfs_update(u);
chmin(dis[cen], min_dep[c[cen]]);
for (int u : edge[cen])
dfs_solve(u);
}
void cal_init() {
for (int i = 1;i <= n;++i) {
dis[i] = inf;
ban[i] = false;
}
}
int cal() {
cal_init();
dfs_solve(1);
return *max_element(dis+1, dis+n+1);
}
int bfs(int s) {
queue<pair<int,int>> q;
vector<bool> vis(n+1);
q.emplace(s, 0);
vis[s] = true;
while (q.size()) {
auto [x, len] = q.front(); q.pop();
if (x != s && c[x] == c[s])
return len;
for (int u : edge[x]) if (!vis[u])
q.emplace(u, len+1), vis[u] = true;
}
return inf;
}
int brute_cal() {
cal_init();
int res = 0;
for (int i = 1;i <= n;++i)
res = max(res, bfs(i));
return res;
}
void brute_solve() {
int res = inf, p;
for (int s = 0;s < 1<<n;++s) {
bool ok = true;
for (int i = 0;i < n;++i)
if (c[i+1] != -1 && c[i+1] != (s>>i&1))
ok = false;
if (!ok) continue;
DE(s);
vector<int> oldc(c, c+n+1);
for (int i = 0;i < n;++i)
c[i+1] = s>>i&1;
int val = brute_cal();
chmin(res, val);
if (res == val) p = s;
copy(AI(oldc), c);
}
if (res > n) {
assert(false);
return cout << -1 << '\n', void();
}
cout << res << '\n';
for (int i = 0;i < n;++i)
cout << (p>>i&1) << " \n"[i+1==n];
}
void solve() {
cin >> n;
clear();
for (int a, b, i = 1;i < n;++i) {
cin >> a >> b;
edge[a].pb(b);
edge[b].pb(a);
}
for (int i = 1;i <= n;++i)
cin >> c[i];
int cn = count(c+1, c+n+1, -1), c0 = count(c+1, c+n+1, 0), c1 = count(c+1, c+n+1, 1);
if ((c0==1) + (c1==1) > cn)
return cout << -1 << '\n', void();
if (n <= 17)
return brute_solve(), void();
if (cn == 0) {
int val = cal();
if (val > n)
return cout << -1 << '\n', void();
cout << val << '\n';
for (int i = 1;i <= n;++i)
cout << c[i] << " \n"[i==n];
}
else
exit(1);
}
int32_t main() {
ios_base::sync_with_stdio(0), cin.tie(0);
int T;
cin >> T;
while (T--)
solve();
}
Compilation message
balancedtree.cpp: In function 'void brute_solve()':
balancedtree.cpp:14:17: warning: statement has no effect [-Wunused-value]
14 | #define DE(...) 0
| ^
balancedtree.cpp:116:3: note: in expansion of macro 'DE'
116 | DE(s);
| ^~
balancedtree.cpp:133:13: warning: 'p' may be used uninitialized in this function [-Wmaybe-uninitialized]
133 | cout << (p>>i&1) << " \n"[i+1==n];
| ~^~~
# |
Verdict |
Execution time |
Memory |
Grader output |
1 |
Correct |
219 ms |
12152 KB |
Output is correct |
2 |
Correct |
933 ms |
12280 KB |
Output is correct |
# |
Verdict |
Execution time |
Memory |
Grader output |
1 |
Correct |
63 ms |
12280 KB |
Output is correct |
2 |
Correct |
198 ms |
14584 KB |
Output is correct |
3 |
Correct |
115 ms |
12792 KB |
Output is correct |
# |
Verdict |
Execution time |
Memory |
Grader output |
1 |
Runtime error |
8 ms |
12032 KB |
Execution failed because the return code was nonzero |
2 |
Runtime error |
50 ms |
15608 KB |
Execution failed because the return code was nonzero |
3 |
Runtime error |
21 ms |
13440 KB |
Execution failed because the return code was nonzero |
# |
Verdict |
Execution time |
Memory |
Grader output |
1 |
Runtime error |
16 ms |
12800 KB |
Execution failed because the return code was nonzero |
2 |
Runtime error |
8 ms |
12160 KB |
Execution failed because the return code was nonzero |
3 |
Runtime error |
8 ms |
12160 KB |
Execution failed because the return code was nonzero |
4 |
Runtime error |
8 ms |
12032 KB |
Execution failed because the return code was nonzero |
5 |
Runtime error |
8 ms |
12160 KB |
Execution failed because the return code was nonzero |
6 |
Runtime error |
10 ms |
12288 KB |
Execution failed because the return code was nonzero |
7 |
Runtime error |
8 ms |
12160 KB |
Execution failed because the return code was nonzero |
8 |
Runtime error |
7 ms |
12032 KB |
Execution failed because the return code was nonzero |
9 |
Runtime error |
8 ms |
12160 KB |
Execution failed because the return code was nonzero |
# |
Verdict |
Execution time |
Memory |
Grader output |
1 |
Correct |
219 ms |
12152 KB |
Output is correct |
2 |
Correct |
933 ms |
12280 KB |
Output is correct |
3 |
Correct |
63 ms |
12280 KB |
Output is correct |
4 |
Correct |
198 ms |
14584 KB |
Output is correct |
5 |
Correct |
115 ms |
12792 KB |
Output is correct |
6 |
Runtime error |
8 ms |
12032 KB |
Execution failed because the return code was nonzero |
7 |
Runtime error |
50 ms |
15608 KB |
Execution failed because the return code was nonzero |
8 |
Runtime error |
21 ms |
13440 KB |
Execution failed because the return code was nonzero |
9 |
Runtime error |
16 ms |
12800 KB |
Execution failed because the return code was nonzero |
10 |
Runtime error |
8 ms |
12160 KB |
Execution failed because the return code was nonzero |
11 |
Runtime error |
8 ms |
12160 KB |
Execution failed because the return code was nonzero |
12 |
Runtime error |
8 ms |
12032 KB |
Execution failed because the return code was nonzero |
13 |
Runtime error |
8 ms |
12160 KB |
Execution failed because the return code was nonzero |
14 |
Runtime error |
10 ms |
12288 KB |
Execution failed because the return code was nonzero |
15 |
Runtime error |
8 ms |
12160 KB |
Execution failed because the return code was nonzero |
16 |
Runtime error |
7 ms |
12032 KB |
Execution failed because the return code was nonzero |
17 |
Runtime error |
8 ms |
12160 KB |
Execution failed because the return code was nonzero |
18 |
Runtime error |
12 ms |
12416 KB |
Execution failed because the return code was nonzero |
19 |
Runtime error |
56 ms |
15608 KB |
Execution failed because the return code was nonzero |
20 |
Runtime error |
8 ms |
12032 KB |
Execution failed because the return code was nonzero |
21 |
Runtime error |
321 ms |
30200 KB |
Execution failed because the return code was nonzero |
22 |
Runtime error |
8 ms |
12160 KB |
Execution failed because the return code was nonzero |