#include <bits/stdc++.h>
#define rep(a,b,c) for(auto a = (b); a != (c); a++)
#define repD(a,b,c) for(auto a = (b); a != (c); a--)
#define repIn(a, b) for(auto& a : (b))
#define repIn2(a, b, c) for(auto& [a, b] : (c))
constexpr bool dbg = 1;
#define DEBUG if constexpr(dbg)
#define DC DEBUG std::cerr
#define eol std::endl
#define ll long long
#define pb push_back
using namespace std;
int n, m, tb;
vector<vector<int>> g;
vector<int> c, t, ans, ans2, depth, maxD, atDepth, cntUniques;
void tAdd(int l, int r, int x) {
if(l > r) return;
DC << " tAdd(" << l << ' ' << r << ' ' << x << ")" << eol;
l += tb; r += tb;
t[l] += x;
if(r != l) t[r] += x;
while(l / 2 != r / 2) {
if(l % 2 == 0) t[l + 1] += x;
if(r % 2 == 1) t[r - 1] += x;
l /= 2; r /= 2;
}
}
int tQ(int i) {
DC << " tQ(" << i << ") = ";
i += tb;
int ret = 0;
while(i) ret += t[i], i /= 2;
DC << ret << eol;
return ret;
}
pair<int, int> findDeepest(int v, int p) {
depth[v] = depth[p] + 1;
maxD[v] = depth[v];
pair<int, int> ret = {depth[v], v};
repIn(u, g[v]) if(u != p) ret = max(ret, findDeepest(u, v)), maxD[v] = max(maxD[v], maxD[u]);
return ret;
}
void dfsAns(int v, int p) {
DC << " Entering " << v << eol;
atDepth[depth[v]] = c[v];
ans[v] = ans[p];
auto myReach = max(1, depth[p] - (maxD[v] - depth[p]));
tAdd(myReach, depth[p] - 1, -1);
repIn(u, g[v]) if(u != p) {
auto reach = max(1, depth[v] - (maxD[u] - depth[v]));
tAdd(reach, depth[p], 1);
}
rep(i, myReach, min(myReach + 2, max(myReach, depth[v]))) if(tQ(i) == 0) {
auto who = atDepth[i];
cntUniques[who]++;
if(cntUniques[who] == 1) ans[v]++;
}
repIn(u, g[v]) if(u != p) {
dfsAns(u, v);
}
DC << " Leaving " << v << eol;
rep(i, myReach, min(myReach + 2, max(myReach, depth[v]))) if(tQ(i) == 0) {
auto who = atDepth[i];
cntUniques[who]--;
}
repIn(u, g[v]) if(u != p) {
auto reach = max(1, depth[v] - (maxD[u] - depth[v]));
tAdd(reach, depth[p], -1);
}
tAdd(myReach, depth[p] - 1, 1);
}
int main() {
ios_base::sync_with_stdio(false); cin.tie(nullptr); cout.tie(nullptr);
cin >> n >> m;
g.resize(n + 1);
c.resize(n + 1);
tb = 1 << (32 - __builtin_clz(n));
t.resize(2 * tb);
ans.resize(n + 1); ans2.resize(n + 1); depth.resize(n + 1); atDepth.resize(n + 1), maxD.resize(n + 1);
cntUniques.resize(m + 1);
rep(i, 1, n) {
int a, b;
cin >> a >> b;
g[a].pb(b);
g[b].pb(a);
}
rep(i, 1, n + 1) cin >> c[i];
auto [_, A] = findDeepest(1, 0);
auto [__, B] = findDeepest(A, 0);
DC << "Root " << A << eol;
dfsAns(A, 0);
swap(ans, ans2);
DC << "Root " << B << eol;
findDeepest(B, 0);
dfsAns(B, 0);
rep(i, 1, n + 1) cout << max(ans[i], ans2[i]) << '\n';
}
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |