This submission is migrated from previous version of oj.uz, which used different machine for grading. This submission may have different result if resubmitted.
#include<bits/stdc++.h>
#define pb push_back
#define ii pair<int,int>
#define all(x) (x).begin(),(x).end()
#define INF 100000000000000000
#define modulo 1000000007
#define mod 998244353
#define int long long int
using namespace std;
struct Tree{
vector<vector<int> >& adj;
vector<vector<int> > anc;
vector<vector<int> > cost;
vector<int> depth;
vector<int> sub;
vector<bool> vis;
int n;
int root;
Tree(vector<vector<int> >& a, int N = 1e6, int r = 1): adj(a){
n = N + 5;
root = r;
anc = vector<vector<int> > (25, vector<int>(n, 0));
depth = vector<int>(n, 0);
vis = vector<bool> (n, false);
sub = vector<int> (n, 1);
init(root, -1);
for(int d = 1; d < 25; d++){
for(int x = 0; x < n; x++){
anc[d][x] = anc[d - 1][anc[d - 1][x]];
}
}
}
void init(int x, int pre){
for(int i = 0; i < adj[x].size(); i++){
int y = adj[x][i];
if(y == pre)continue;
anc[0][y] = x;
depth[y] = depth[x] + 1;
init(y, x);
sub[x] += sub[y];
}
}
int LCA(int x, int y){
if(depth[x] < depth[y])swap(x,y);
while(depth[x] != depth[y]){
int diff = depth[x] - depth[y];
diff = log2(diff&-diff);
x = anc[diff][x];
}
if(x == y) return x;
for(int i = 24; i >= 0; i--){
if(anc[i][x] != anc[i][y]){
x = anc[i][x];
y = anc[i][y];
}
}
return anc[0][x];
}
int up(int x, int k){
while(k){
int diff = log2(k&-k);
k -= k & -k;
x = anc[diff][x];
}
return x;
}
int dist(int x, int y){
int l = LCA(x, y);
return depth[x] + depth[y] - 2 * depth[l];
}
};
vector<vector<int> > adj(500005);
vector<int> val(500005);
vector<int> comp[500005];
vector<int> up(500005);
vector<int> sub(500005);
unordered_map<int,int> edge;
vector<bool> virt(500005, false);
int num(int x, int y){
return edge[x * 1e6 + y];
}
void dfs(int x, int pre, vector<int>& depth){
for(auto y : adj[x]){
if(y == pre) continue;
dfs(y, x, depth);
if(depth[up[x]] > depth[up[y]])
up[x] = up[y];
}
if(pre != 0 && up[x] != x){
virt[num(x, pre)] = true;
}
}
void get(int x, int pre, int& ans, int root){
sub[x] = 0;
if(pre != 0 && virt[num(x, pre)] == false) sub[x] = 1;
for(auto y : adj[x]){
if(y == pre)continue;
get(y, x, ans, root);
sub[x] += sub[y];
}
if(pre != 0 && pre != root && virt[num(x, pre)] == false && sub[x] == 1) ans++;
}
int32_t main(){
ios_base::sync_with_stdio(false);
cin.tie(NULL);
int n, k;
cin >> n >> k;
if(n == 1){
cout << 0;
return 0;
}
for(int i = 0; i < n - 1; i++) {
int x, y;
cin >> x >> y;
adj[x].pb(y);
adj[y].pb(x);
edge[x * 1e6 + y] = i;
edge[y * 1e6 + x] = i;
}
for(int i = 1; i <= n; i++){
cin >> val[i];
comp[val[i]].pb(i);
up[i] = i;
}
Tree A(adj);
for(int i = 0; i < 500005; i++){
if(comp[i].empty()) continue;
int l = comp[i][0];
for(auto x : comp[i]){
l = A.LCA(x, l);
}
for(auto x : comp[i]){
up[x] = l;
}
}
dfs(1, 0, A.depth);
int root = 0;
for(auto itr = edge.begin(); itr != edge.end(); itr++) {
int x = itr->first / 1e6;
int y = itr->first % 1000000;
int w = itr->second;
if(virt[w]) continue;
if(A.depth[x] < A.depth[y]) swap(x, y);
if((root == 0 || A.depth[x] > A.depth[root])) root = x;
}
int ans = 0;
get(root, 0, ans, root);
if(root) ans++;
cout << (ans + 1) / 2;
}
Compilation message (stderr)
mergers.cpp: In member function 'void Tree::init(long long int, long long int)':
mergers.cpp:45:30: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
for(int i = 0; i < adj[x].size(); i++){
~~^~~~~~~~~~~~~~~
# | 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... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |