#include <algorithm>
#include <iostream>
#include <vector>
using namespace std;
const int N = 200000;
vector<int> ej[N], ej_[N];
int aa[N], dds[N], ddt[N], *dd, d_, i_;
int hh[N], ii_[N], cc_[N];
int ii[N], pp[N], qq[N], rr[N], n_;
int kk[N], cc[N];
void dfs_dd(int p, int i, int d) {
if (d_ < d)
d_ = d, i_ = i;
dd[i] = d++;
for (int j : ej[i])
if (j != p)
dfs_dd(i, j, d);
}
int dfs_hh(int p, int i) {
int h = 0;
for (int j : ej[i])
if (j != p)
h = max(h, dfs_hh(i, j) + 1);
return hh[i] = h;
}
int search(int i_, int d) {
while (i_ && dds[ii[i_]] >= d)
i_ = qq[i_] && dds[ii[qq[i_]]] >= d ? qq[i_] : pp[i_];
return i_;
}
void dfs(int p, int i, int i_) {
if (dds[i] >= ddt[i])
ii_[i] = search(i_, dds[i] - hh[i]);
int h0 = 0, h1 = 0;
for (int j : ej[i])
if (j != p) {
int h = hh[j] + 1;
if (h0 < h)
h1 = h0, h0 = h;
else if (h1 < h)
h1 = h;
}
for (int j : ej[i])
if (j != p) {
int p_ = search(i_, dds[i] - (hh[j] + 1 == h0 ? h1 : h0));
ii[n_] = i, ej_[pp[n_] = p_].push_back(n_);
if (rr[p_] == rr[qq[p_]])
qq[n_] = qq[qq[p_]], rr[n_] = rr[p_] + 1;
else
qq[n_] = p_, rr[n_] = 0;
dfs(i, j, n_++);
}
}
void dfs_kk(int i) {
static int c;
cc[i] = c += i && !kk[aa[ii[i]]]++;
for (int j : ej_[i])
dfs_kk(j);
c -= i && !--kk[aa[ii[i]]];
}
int main() {
ios_base::sync_with_stdio(false), cin.tie(NULL);
int n, a_; cin >> n >> a_;
for (int h = 0; h < n - 1; h++) {
int i, j; cin >> i >> j, i--, j--;
ej[i].push_back(j), ej[j].push_back(i);
}
for (int i = 0; i < n; i++)
cin >> aa[i], aa[i]--;
dd = dds, d_ = -1, dfs_dd(-1, 0, 0); int s = i_;
dd = dds, d_ = -1, dfs_dd(-1, s, 0); int t = i_;
dd = ddt, dfs_dd(-1, t, 0);
for (int z = 0; z < 2; z++) {
dfs_hh(-1, s);
n_ = 1, dfs(-1, s, 0);
dfs_kk(0);
for (int i = 0; i < n; i++)
if (dds[i] >= ddt[i])
cc_[i] = cc[ii_[i]];
swap(s, t);
for (int i = 0; i < n; i++)
swap(dds[i], ddt[i]), ej_[i].clear();
}
for (int i = 0; i < n; i++)
cout << cc_[i] << '\n';
return 0;
}
# | 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... |