#include <bits/stdc++.h>
#define all(i) (i).begin(), (i).end()
using namespace std;
template<typename T1>
void Debug(bool _split, T1 i) {
if (_split)
cerr << ", ";
cerr << i;
}
template<typename T1, typename ...T2>
void Debug(bool _split, T1 i, T2 ...j) {
if (_split)
cerr << ", ";
cerr << i;
Debug(true, j...);
}
#define debug(args...) cout << "Line" << __LINE__ << " : [" << #args << "] is [", Debug(false, args), cerr << "]" << endl;
typedef pair<int, int> pi;
typedef long long ll;
const int inf = 0x3f3f3f3f, lg = 20;
const ll mod = 1e9 + 7, INF = 0x3f3f3f3f;
const int N = 2e5;
vector<int> g[N];
int pa[N], c[N], at[N], cnt[N], ans[N];
pair<pi, pi> mx[N];
vector<pi> vv;
int now;
pi dfs1(int i, int p) {
pa[i] = p;
pi res(0, i);
for (int j : g[i])
if (j != p) {
pi tmp = dfs1(j, i);
++tmp.first;
res = max(res, tmp);
}
return res;
}
int dfs2(int i, int p, int _v) {
mx[i].first = mx[i].second = pi(0, i);
for (int j : g[i])
if (j != p) {
pi tmp(dfs2(j, i, _v - 1) + 1, j);
if (tmp.first > mx[i].first.first)
mx[i].second = mx[i].first, mx[i].first = tmp;
else if (tmp.first > mx[i].second.first)
mx[i].second = tmp;
}
return mx[i].first.first;
}
void dfs3(int i, int p, int _v) {
if (at[_v] == -2)
at[_v] = c[i];
else
at[_v] = -1;
for (int j : g[i])
if (j != p)
dfs3(j, i, _v + 1);
}
void add(int x, int y) {
vv.emplace_back(x, y);
if (!cnt[y])
++now;
++cnt[y];
}
void del(int x = inf, int y = -1) {
if (~y && (vv.empty() || vv.back() != pi(x, y)))
return;
y = vv.back().second;
vv.pop_back();
--cnt[y];
if (!cnt[y])
--now;
}
void dfs4(int i, int p, int _v) {
int tmp;
if (mx[i].first.second != i) {
if (mx[i].second.first > 0)
tmp = _v + mx[i].second.first;
else
tmp = _v;
while (vv.back().first <= tmp)
del();
add(_v, c[i]);
dfs4(mx[i].first.second, i, _v - 1);
del(_v, c[i]);
}
if (mx[i].first.first > 0)
tmp = _v + mx[i].first.first;
else
tmp = _v;
while (vv.back().first <= tmp)
del();
ans[i] = now;
add(_v, c[i]);
for (int j : g[i])
if (j != p && j != mx[i].first.second)
dfs4(j, i, _v - 1);
del(_v, c[i]);
}
signed main() {
ios::sync_with_stdio(0), cin.tie(0);
int n, _;
cin >> n >> _;
for (int i = 0; i < n - 1; ++i) {
int x, y;
cin >> x >> y;
--x, --y;
g[x].push_back(y), g[y].push_back(x);
}
for (int i = 0; i < n; ++i)
cin >> c[i], --c[i];
int x = dfs1(0, -1).second;
int y = dfs1(x, -1).second;
vector<int> dm;
for (int i = y; ~i; i = pa[i])
dm.push_back(i);
x = dm.size() - 1 >> 1;//2k -> k-1, 2k+1 -> k
dfs2(dm[x], dm[x + 1], 0);
fill(at, at + n, -2);
dfs3(dm[x + 1], dm[x], 1);
fill(cnt, cnt + n, 0);
now = 0;
vv.emplace_back(inf, -1);
for (int i = n - 1; i; --i)
if (at[i] >= 0) {
add(i, at[i]);
}
dfs4(dm[x], dm[x + 1], 0);
x = dm.size() >> 1;//2k -> k, 2k+1 -> k
dfs2(dm[x], dm[x - 1], 0);
fill(at, at + n, -2);
dfs3(dm[x - 1], dm[x], 1);
fill(cnt, cnt + n, 0);
now = 0;
vv.clear();
vv.emplace_back(inf, -1);
for (int i = n - 1; i; --i)
if (at[i] >= 0) {
add(i, at[i]);
}
dfs4(dm[x], dm[x - 1], 0);
for (int i = 0; i < n; ++i)
cout << ans[i] << '\n';
}
Compilation message
joi2019_ho_t5.cpp: In function 'int main()':
joi2019_ho_t5.cpp:142:19: warning: suggest parentheses around '-' inside '>>' [-Wparentheses]
142 | x = dm.size() - 1 >> 1;//2k -> k-1, 2k+1 -> k
| ~~~~~~~~~~^~~
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
3 ms |
4940 KB |
Output is correct |
| 2 |
Incorrect |
4 ms |
5156 KB |
Output isn't correct |
| 3 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Incorrect |
100 ms |
13252 KB |
Output isn't correct |
| 2 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Incorrect |
170 ms |
16320 KB |
Output isn't correct |
| 2 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
3 ms |
4940 KB |
Output is correct |
| 2 |
Incorrect |
4 ms |
5156 KB |
Output isn't correct |
| 3 |
Halted |
0 ms |
0 KB |
- |
