#include <bits/stdc++.h>
#include <unordered_map>
#pragma GCC optimize("O3")
#pragma GCC optimize("Ofast")
#pragma GCC optimize("unroll-loops")
#pragma GCC target("avx2")
using namespace std;
typedef long long ll;
typedef pair<ll, ll> pll;
#define MAX 201010
#define MAXS 20
#define INF 1000000000000000001
#define bb ' '
#define ln '\n'
struct segtree {
ll N;
ll s;
vector<ll> tree, l, r;
void update(ll x, ll a) {
x += s - 1;
tree[x] = a;
x /= 2;
while (x) tree[x] = max(tree[x * 2], tree[x * 2 + 1]), x /= 2;
}
ll query(ll low, ll high, ll loc = 1) {
if (l[loc] == low && r[loc] == high) return tree[loc];
if (r[loc * 2] >= high) return query(low, high, loc * 2);
if (l[loc * 2 + 1] <= low) return query(low, high, loc * 2 + 1);
return max(query(low, r[loc * 2], loc * 2), query(l[loc * 2 + 1], high, loc * 2 + 1));
}
void init(ll x = 1) {
if (x >= s) {
l[x] = r[x] = x - s + 1;
return;
}
init(x * 2);
init(x * 2 + 1);
l[x] = l[x * 2];
r[x] = r[x * 2 + 1];
}
segtree(ll n) {
N = n;
s = 1LL << (ll)ceil(log2(N));
tree.resize(2 * s + 1);
l.resize(2 * s + 1);
r.resize(2 * s + 1);
init();
}
};
vector<ll> adj[MAX], sav[MAX], rev[MAX];
ll dir[MAX], ddir[MAX], dddir[MAX];
vector<vector<ll>> chain;
vector<set<ll>> subtree, revtree;
vector<segtree> chainseg;
ll C[MAX], depth[MAX], mxdepv[MAX], prtval[MAX];
pll range[MAX];
ll sp[MAX][MAXS];
ll cnt;
ll num[MAX];
pll cnum[MAX];
ll ans[MAX];
ll init(ll x, ll p = 0, ll d = 0) {
sp[x][0] = p;
ll i;
for (i = 1; i < MAXS; i++) sp[x][i] = sp[sp[x][i - 1]][i - 1];
depth[x] = d;
ll mx = 0;
ll sum = 1;
sav[x].resize(adj[x].size());
rev[x].resize(adj[x].size());
for (auto v : adj[x]) {
if (v == p) continue;
sum += init(v, x, d + 1);
}
return num[x] = sum;
}
void calc(ll x, ll p = 0) {
ll i;
ll mx = 0;
mxdepv[x] = x;
for (i = 0; i < adj[x].size(); i++) {
if (adj[x][i] == p) continue;
calc(adj[x][i], x);
if (mx < depth[mxdepv[adj[x][i]]]) mx = depth[mxdepv[adj[x][i]]], mxdepv[x] = mxdepv[adj[x][i]];
sav[x][i] = mxdepv[adj[x][i]];
}
ll cnt = 0;
ll vv, vvv;
vv = vvv = -1;
ll nv = -1;
ll nmx = 0;
for (auto v : adj[x]) {
if (v == p) continue;
if (depth[mxdepv[v]] == mx) cnt++, vvv = vv, vv = v;
else if (depth[mxdepv[v]] > nmx) nmx = depth[mxdepv[v]], nv = v;
}
if (cnt >= 2) {
for (i = 0; i < adj[x].size(); i++) {
if (adj[x][i] != p) {
rev[x][i] = (adj[x][i] == vv ? vvv : vv);
}
}
}
else {
for (i = 0; i < adj[x].size(); i++) {
if (adj[x][i] != p) {
rev[x][i] = (adj[x][i] == vv ? nv : vv);
}
}
}
for (i = 0; i < adj[x].size(); i++) {
if (adj[x][i] == p) continue;
if (rev[x][i] == -1) rev[x][i] = x;
else rev[x][i] = mxdepv[rev[x][i]];
}
}
void make_chain(ll x, ll p = 0) {
ll mx, mv;
mx = mv = 0;
chain[cnt].push_back(x);
cnum[x] = { cnt, chain[cnt].size() - 1 };
for (auto v : adj[x]) {
if (v == p) continue;
if (mx < num[v]) mx = num[v], mv = v;
}
if (mv) make_chain(mv, x);
for (auto v : adj[x]) {
if (v == p || v == mv) continue;
cnt++;
chain.push_back(vector<ll>());
make_chain(v, x);
}
}
void make_tree() {
ll i;
for (i = 0; i < chain.size(); i++) chainseg.push_back(segtree(chain[i].size()));
}
void update(ll v, ll x) {
chainseg[cnum[v].first].update(cnum[v].second + 1, x);
}
//1을 루트로 하는 LCA
ll lca(ll u, ll v) {
if (depth[u] != depth[v]) {
if (depth[u] < depth[v]) swap(u, v);
ll i;
for (i = MAXS - 1; i >= 0; i--) if (depth[sp[u][i]] >= depth[v]) u = sp[u][i];
}
if (u == v) return u;
ll i;
for (i = MAXS - 1; i >= 0; i--) if (sp[u][i] != sp[v][i]) u = sp[u][i], v = sp[v][i];
return sp[v][0];
}
//HLD query
ll mxval(ll u, ll v) {
ll ans = 0;
ll l = lca(u, v);
while (cnum[u].first != cnum[l].first) ans = max(ans, chainseg[cnum[u].first].query(1, cnum[u].second + 1)), u = sp[chain[cnum[u].first][0]][0];
while (cnum[v].first != cnum[l].first) ans = max(ans, chainseg[cnum[v].first].query(1, cnum[v].second + 1)), v = sp[chain[cnum[v].first][0]][0];
ans = max(ans, chainseg[cnum[l].first].query(cnum[l].second + 1, cnum[u].second + 1));
ans = max(ans, chainseg[cnum[l].first].query(cnum[l].second + 1, cnum[v].second + 1));
return ans;
}
//두 정점 사이 거리
ll dis(ll u, ll v) { return depth[u] + depth[v] - 2 * depth[lca(u, v)]; }
//r이 루트, v의 x번째 부모
ll prtx(ll r, ll v, ll x) {
if (x == 0) return v;
ll rv = dis(r, v);
if (rv < x) return 0;
ll l = lca(r, v);
if (dis(l, v) < x) {
ll d = rv - x;
ll i;
for (i = MAXS - 1; i >= 0; i--) if (d - (1 << i) >= 0) d -= (1 << i), r = sp[r][i];
return r;
}
else {
ll i;
for (i = MAXS - 1; i >= 0; i--) if (x - (1 << i) >= 0) x -= (1 << i), v = sp[v][i];
return v;
}
}
ll getfar(ll v, ll ban) {
if (dir[v] != ban) return dir[v];
return ddir[v];
}
ll getfar(ll v, ll ban1, ll ban2) {
if (ban1 > ban2) swap(ban1, ban2);
if (ban2 == -1) return dir[v];
if (ban1 == -1) return getfar(v, ban2);
if (dir[v] != ban1 && dir[v] != ban2) return dir[v];
if (ddir[v] != ban1 && ddir[v] != ban2) return ddir[v];
return dddir[v];
}
ll getind(vector<ll>& v, ll c) {
return lower_bound(v.begin(), v.end(), c) - v.begin();
}
//r1 : previous root
void prop(ll r1, ll r2) {
if (ddir[r2] == -1) update(r2, dis(sav[r2][dir[r2]], r2));
else update(r2, dis(r2, sav[r2][dir[r2]]) + dis(r2, sav[r2][ddir[r2]]));
ll ind = getind(adj[r1], r2);
ll f1 = getfar(r1, ind);
ll f2 = getfar(r1, ind, f1);
if (f1 == -1) update(r1, 0);
else if (f2 == -1) update(r1, dis(r1, sav[r1][f1]));
else update(r1, dis(r1, sav[r1][f1]) + dis(r1, sav[r1][f2]));
}
void dfs(ll x, ll p = 0) {
ll i;
for (i = 0; i < adj[x].size(); i++) {
ll fardir = getfar(adj[x][i], getind(adj[adj[x][i]], x));
ll farv = sav[x][i];
ll fardis = dis(x, farv);
ll res = mxval(farv, adj[x][i]);
if (res >= fardis) continue;
ll xx = (fardis - 1) / 2;
ll root = prtx(x, farv, xx);
if (lca(root, x) == root) revtree[prtx(x, root, 1)].insert(C[x]);
else subtree[root].insert(C[x]);
}
for (auto v : adj[x]) {
if (v == p) continue;
prop(x, v);
dfs(v, x);
prop(v, x);
}
}
unordered_map<ll, ll> mp;
void getans(ll x, ll p = 0) {
ll i;
for (auto c : subtree[x]) mp[c]++;
vector<ll> vv;
for (auto c : revtree[x]) {
mp[c]--;
if (!mp[c]) vv.push_back(c);
}
for (auto x : vv) mp.erase(x);
vv.clear();
ans[x] = mp.size();
for (auto v : adj[x]) {
if (v == p) continue;
getans(v, x);
}
for (auto c : revtree[x]) mp[c]++;
for (auto c : subtree[x]) {
mp[c]--;
if (!mp[c]) vv.push_back(c);
}
for (auto x : vv) mp.erase(x);
}
void calcp(ll x, ll p = 0) {
ll i, v;
if (x != 1) {
if (p == 1) sav[x][getind(adj[x], p)] = rev[p][getind(adj[p], x)];
else {
ll v1 = rev[p][getind(adj[p], x)];
ll v2 = sav[p][getind(adj[p], sp[p][0])];
if (v1 > v2) swap(v1, v2);
if (dis(x, v1) >= dis(x, v2)) sav[x][getind(adj[x], p)] = v1;
else sav[x][getind(adj[x], p)] = v2;
}
}
for (auto v : adj[x]) if (v != p) calcp(v, x);
}
signed main() {
ios::sync_with_stdio(false), cin.tie(0);
depth[0] = -1;
ll N, M;
cin >> N >> M;
ll i, j;
ll a, b;
for (i = 1; i < N; i++) cin >> a >> b, adj[a].push_back(b), adj[b].push_back(a);
for (i = 1; i <= N; i++) cin >> C[i];
for (i = 1; i <= N; i++) sort(adj[i].begin(), adj[i].end());
init(1);
calc(1);
calcp(1);
cnt = 0;
chain.push_back(vector<ll>());
make_chain(1);
make_tree();
return 0;
for (i = 1; i <= N; i++) {
ll mx = 0, mv = 0;
dir[i] = ddir[i] = dddir[i] = -1;
for (j = 0; j < adj[i].size(); j++) if (mx < dis(i, sav[i][j])) mx = dis(i, sav[i][j]), dir[i] = j;
mx = 0;
for (j = 0; j < adj[i].size(); j++) {
if (j == dir[i]) continue;
if (mx < dis(i, sav[i][j])) mx = dis(i, sav[i][j]), ddir[i] = j;
}
mx = 0;
for (j = 0; j < adj[i].size(); j++) {
if (j == dir[i] || j == ddir[i]) continue;
if (mx < dis(i, sav[i][j])) mx = dis(i, sav[i][j]), dddir[i] = j;
}
if (i != 1) {
ll p = getind(adj[i], sp[i][0]);
ll r = getfar(i, p);
ll rr = getfar(i, p, r);
ll xx = 0;
if (r != -1) xx += dis(i, sav[i][r]);
if (rr != -1) xx += dis(i, sav[i][rr]);
update(i, xx);
}
else {
ll dd;
dd = ddir[1];
if (dd == -1) update(1, depth[sav[1][dir[1]]]);
else update(1, depth[sav[1][dd]] + depth[sav[1][dir[1]]]);
}
}
subtree.resize(N + 1);
revtree.resize(N + 1);
dfs(1);
for (i = 1; i <= N; i++) for (auto c : revtree[i]) mp[c]++;
getans(1);
for (i = 1; i <= N; i++) cout << ans[i] << ln;
}
Compilation message
joi2019_ho_t5.cpp: In function 'll init(ll, ll, ll)':
joi2019_ho_t5.cpp:67:5: warning: unused variable 'mx' [-Wunused-variable]
67 | ll mx = 0;
| ^~
joi2019_ho_t5.cpp: In function 'void calc(ll, ll)':
joi2019_ho_t5.cpp:81:16: warning: comparison of integer expressions of different signedness: 'll' {aka 'long long int'} and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
81 | for (i = 0; i < adj[x].size(); i++) {
| ~~^~~~~~~~~~~~~~~
joi2019_ho_t5.cpp:98:17: warning: comparison of integer expressions of different signedness: 'll' {aka 'long long int'} and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
98 | for (i = 0; i < adj[x].size(); i++) {
| ~~^~~~~~~~~~~~~~~
joi2019_ho_t5.cpp:105:17: warning: comparison of integer expressions of different signedness: 'll' {aka 'long long int'} and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
105 | for (i = 0; i < adj[x].size(); i++) {
| ~~^~~~~~~~~~~~~~~
joi2019_ho_t5.cpp:111:16: warning: comparison of integer expressions of different signedness: 'll' {aka 'long long int'} and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
111 | for (i = 0; i < adj[x].size(); i++) {
| ~~^~~~~~~~~~~~~~~
joi2019_ho_t5.cpp: In function 'void make_tree()':
joi2019_ho_t5.cpp:136:16: warning: comparison of integer expressions of different signedness: 'll' {aka 'long long int'} and 'std::vector<std::vector<long long int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
136 | for (i = 0; i < chain.size(); i++) chainseg.push_back(segtree(chain[i].size()));
| ~~^~~~~~~~~~~~~~
joi2019_ho_t5.cpp: In function 'void dfs(ll, ll)':
joi2019_ho_t5.cpp:211:16: warning: comparison of integer expressions of different signedness: 'll' {aka 'long long int'} and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
211 | for (i = 0; i < adj[x].size(); i++) {
| ~~^~~~~~~~~~~~~~~
joi2019_ho_t5.cpp:212:6: warning: unused variable 'fardir' [-Wunused-variable]
212 | ll fardir = getfar(adj[x][i], getind(adj[adj[x][i]], x));
| ^~~~~~
joi2019_ho_t5.cpp: In function 'void getans(ll, ll)':
joi2019_ho_t5.cpp:231:5: warning: unused variable 'i' [-Wunused-variable]
231 | ll i;
| ^
joi2019_ho_t5.cpp: In function 'void calcp(ll, ll)':
joi2019_ho_t5.cpp:253:5: warning: unused variable 'i' [-Wunused-variable]
253 | ll i, v;
| ^
joi2019_ho_t5.cpp:253:8: warning: unused variable 'v' [-Wunused-variable]
253 | ll i, v;
| ^
joi2019_ho_t5.cpp: In function 'int main()':
joi2019_ho_t5.cpp:287:17: warning: comparison of integer expressions of different signedness: 'll' {aka 'long long int'} and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
287 | for (j = 0; j < adj[i].size(); j++) if (mx < dis(i, sav[i][j])) mx = dis(i, sav[i][j]), dir[i] = j;
| ~~^~~~~~~~~~~~~~~
joi2019_ho_t5.cpp:289:17: warning: comparison of integer expressions of different signedness: 'll' {aka 'long long int'} and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
289 | for (j = 0; j < adj[i].size(); j++) {
| ~~^~~~~~~~~~~~~~~
joi2019_ho_t5.cpp:294:17: warning: comparison of integer expressions of different signedness: 'll' {aka 'long long int'} and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
294 | for (j = 0; j < adj[i].size(); j++) {
| ~~^~~~~~~~~~~~~~~
joi2019_ho_t5.cpp:285:14: warning: unused variable 'mv' [-Wunused-variable]
285 | ll mx = 0, mv = 0;
| ^~
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Incorrect |
9 ms |
14412 KB |
Output isn't correct |
| 2 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Incorrect |
369 ms |
65744 KB |
Output isn't correct |
| 2 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Incorrect |
502 ms |
84588 KB |
Output isn't correct |
| 2 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Incorrect |
9 ms |
14412 KB |
Output isn't correct |
| 2 |
Halted |
0 ms |
0 KB |
- |
