#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
using namespace std;
using namespace __gnu_pbds;
#define int long long
#define mp make_pair
#define mt make_tuple
#define pb push_back
#define ppb pop_back
#define eb emplace_back
#define g0(a) get<0>(a)
#define g1(a) get<1>(a)
#define g2(a) get<2>(a)
#define g3(a) get<3>(a)
#define g4(a) get<4>(a)
mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());
using db = double;
using ll = long long;
using ld = long double;
using ii = pair<int, int>;
using iii = tuple<int, int, int>;
using iiii = tuple<int, int, int, int>;
using iiiii = tuple<int, int, int, int, int>;
template<class key, class value = null_type, class cmp = less<key> >
using ordered_set = tree<key, value, cmp, rb_tree_tag, tree_order_statistics_node_update>;
int N, M, lon[200005], ans_st[200005], ans_ed[200005], lon_st[200005], lon_ed[200005], dist_st[200005], dist_ed[200005], C[200005];
ii m = mp(-1, -1);
vector<int> adj[200005];
stack<int> rep[200005];
struct node {
node *left, *right;
int S, E, pv;
ii val;
void combine() {
if (S == E) return;
if (left->val.first == right->val.first) val = mp(left->val.first, left->val.second + right->val.second);
else val = min(left->val, right->val);
}
node(int _s, int _e) : S(_s), E(_e), pv(0) {
if (S == E) {
val = mp(0, 1);
return;
}
int M = (S + E) >> 1;
left = new node(S, M);
right = new node(M + 1, E);
combine();
}
void prop() {
if (S == E || !pv) return;
left->val.first += pv;
left->pv += pv;
right->val.first += pv;
right->pv += pv;
pv = 0;
}
void upd(int l, int r, int v) {
if (l > E || r < S) return;
if (l <= S && E <= r) {
val.first += v;
pv += v;
return;
}
prop();
left->upd(l, r, v);
right->upd(l, r, v);
combine();
}
ii qry(int l, int r) {
if (l > E || r < S) return mp((int)1e16, 0);
if (l <= S && E <= r) return val;
prop();
auto lq = left->qry(l, r), rq = right->qry(l, r);
if (lq.first == rq.first) return mp(lq.first, lq.second + rq.second);
else return min(lq, rq);
}
} *root;
void dfs(int n, int e = -1, int d = 0) {
m = max(m, mp(d, n));
for (auto u : adj[n]) if (u != e) {
dfs(u, n, d + 1);
}
}
void init(int n, int lon[], int dist[], int e = -1) {
for (auto u : adj[n]) if (u != e) {
dist[u] = dist[n] + 1;
init(u, lon, dist, n);
lon[n] = max(lon[n], lon[u]);
}
lon[n]++;
}
void get_ans(int n, int ans[], int lon[], int e = -1, int d = 0) {
bool pushed = 0;
if (rep[C[n]].empty() || root->qry(rep[C[n]].top(), rep[C[n]].top()).first) {
rep[C[n]].push(d);
pushed = 1;
} else {
root->upd(d, d, 1);
}
root->upd(d - lon[n] + 1, d - 1, 1);
auto tmp = root->qry(0, d);
ans[n] = (tmp.first ? 0ll : tmp.second - pushed);
root->upd(d - lon[n] + 1, d - 1, -1);
multiset<int> m;
for (auto u : adj[n]) if (u != e) {
m.insert(lon[u]);
}
for (auto u : adj[n]) if (u != e) {
m.erase(m.find(lon[u]));
int x = (m.empty() ? 0ll : *m.rbegin());
root->upd(d - x, d - 1, 1);
get_ans(u, ans, lon, n, d + 1);
root->upd(d - x, d - 1, -1);
m.insert(lon[u]);
}
if (pushed) {
assert(!rep[C[n]].empty() && rep[C[n]].top() == d);
rep[C[n]].pop();
} else {
root->upd(d, d, -1);
}
}
main() {
ios::sync_with_stdio(0);
cin.tie(0);
cin >> N >> M;
for (int i = 1, u, v; i < N; i++) {
cin >> u >> v;
adj[u].pb(v);
adj[v].pb(u);
}
for (int i = 1; i <= N; i++) {
cin >> C[i];
}
dfs(1);
int st = m.second;
m = mp(-1, -1);
dfs(st);
int ed = m.second;
root = new node(0, N - 1);
init(st, lon_st, dist_st);
get_ans(st, ans_st, lon_st);
root = new node(0, N - 1);
init(ed, lon_ed, dist_ed);
for (int i = 1; i <= M; i++) {
while (!rep[i].empty()) rep[i].pop();
}
get_ans(ed, ans_ed, lon_ed);
for (int i = 1; i <= N; i++) {
if (dist_st[i] <= dist_ed[i]) {
cout << ans_ed[i] << '\n';
} else {
cout << ans_st[i] << '\n';
}
}
}
Compilation message
joi2019_ho_t5.cpp:130:1: warning: ISO C++ forbids declaration of 'main' with no type [-Wreturn-type]
130 | main() {
| ^~~~
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
88 ms |
139628 KB |
Output is correct |
| 2 |
Incorrect |
92 ms |
140308 KB |
Output isn't correct |
| 3 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Incorrect |
582 ms |
179520 KB |
Output isn't correct |
| 2 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
734 ms |
194860 KB |
Output is correct |
| 2 |
Correct |
1199 ms |
245156 KB |
Output is correct |
| 3 |
Correct |
230 ms |
153860 KB |
Output is correct |
| 4 |
Correct |
995 ms |
209588 KB |
Output is correct |
| 5 |
Correct |
1227 ms |
250136 KB |
Output is correct |
| 6 |
Correct |
1188 ms |
229344 KB |
Output is correct |
| 7 |
Correct |
913 ms |
210892 KB |
Output is correct |
| 8 |
Correct |
953 ms |
219848 KB |
Output is correct |
| 9 |
Correct |
1034 ms |
216832 KB |
Output is correct |
| 10 |
Correct |
981 ms |
213896 KB |
Output is correct |
| 11 |
Correct |
864 ms |
214176 KB |
Output is correct |
| 12 |
Correct |
1123 ms |
243912 KB |
Output is correct |
| 13 |
Correct |
1037 ms |
230472 KB |
Output is correct |
| 14 |
Correct |
1066 ms |
229556 KB |
Output is correct |
| 15 |
Correct |
855 ms |
219396 KB |
Output is correct |
| 16 |
Correct |
1048 ms |
239712 KB |
Output is correct |
| 17 |
Correct |
1097 ms |
231804 KB |
Output is correct |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
88 ms |
139628 KB |
Output is correct |
| 2 |
Incorrect |
92 ms |
140308 KB |
Output isn't correct |
| 3 |
Halted |
0 ms |
0 KB |
- |
