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>
using namespace std;
typedef long long ll;
typedef long double ld;
typedef pair<ll,ll> pl;
typedef pair<int,int> pii;
typedef tuple<int,int,int> tt;
#define all(a) a.begin(), a.end()
#define filter(a) a.erase(unique(all(a)), a.end())
const int mn = 3e5 + 5;
int up[mn][19], depth[mn], curLCA[mn], deg[mn];
vector<int> state[mn], adj[mn];
bool mark[mn], vis[mn];
vector<pii> edge;
void preDfs (int u, int p, int d) {
    up[u][0] = p, depth[u] = d;
    for (int i = 1; i < 19; i++)
        up[u][i] = up[up[u][i - 1]][i - 1];
    for (int v : adj[u])
        if (v != p) preDfs(v, u, d + 1);
}
int goUp (int a, int k) {
    for (int i = 0; i < 19; i++)
        if (k & (1 << i)) a = up[a][i];
    return a;
}
int lca (int a, int b) {
    if (depth[b] > depth[a]) swap(a, b);
    a = goUp(a, depth[a] - depth[b]);
    if (a == b) return a;
    for (int i = 18; i >= 0; i--)
        if (up[a][i] != up[b][i])
            a = up[a][i], b = up[b][i];
    return up[a][0];
}
struct DSU {
    vector<int> lab;
    DSU (int sz) : lab(sz + 1, -1) {}
    int get (int u) {
        if (lab[u] < 0) return u;
        return lab[u] = get(lab[u]);
    }
    void unite (int a, int b) {
        a = get(a), b = get(b);
        if (a == b) return;
        if (-lab[a] < -lab[b]) swap(a, b);
        lab[a] += lab[b], lab[b] = a;
    }
} dsu(mn);
void dfs (int u, int p) {
    for (int v : adj[u]) {
        if (v == p) continue;
        dfs(v, u);
        curLCA[u] = lca(curLCA[u], curLCA[v]);
    }
    if (curLCA[u] == u && u != p) {
        // for edge that separates the tree into groups
        // create an edge in the new graph
        edge.push_back({u, p});
    }
    else if (u != p) {
        // compress nodes for other edges
        dsu.unite(u, p);
    }
}
int main()
{
    ios::sync_with_stdio(0);
    cin.tie(0);
    int n, k; cin >> n >> k;
    for (int i = 1; i < n; i++) {
        int a, b; cin >> a >> b;
        adj[a].push_back(b);
        adj[b].push_back(a);
    }
    preDfs(1, 1, 1);
    for (int i = 1; i <= n; i++) {
        int j; cin >> j;
        state[j].push_back(i);
    }
    for (int i = 1; i <= k; i++) {
        int cur = 0;
        for (int u : state[i]) {
            if (!cur) cur = u;
            else cur = lca(cur, u);
        }
        for (int u : state[i]) curLCA[u] = cur;
    }
    dfs(1, 1);
    for (pii it : edge) {
        int a, b; tie(a, b) = it;
        a = dsu.get(a), b = dsu.get(b);
        deg[a]++, deg[b]++;
    }
    int ans = 0;
    for (int i = 1; i <= n; i++) {
        int cur = dsu.get(i);
        if (vis[cur]) continue;
        vis[cur] = 1, ans += (deg[cur] == 1 ? 1 : 0);
        // greedy: every pair of leaves (in the new graph)
        // can be merged to satisfy the constraints
    }
    cout << (ans + 1) / 2;
    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... | 
| # | Verdict  | Execution time | Memory | Grader output | 
|---|
| Fetching results... |