Submission #362356

#TimeUsernameProblemLanguageResultExecution timeMemory
362356Mahdi_ShokoufiMergers (JOI19_mergers)C++17
100 / 100
1206 ms137740 KiB
//In the name of Allah
#include <bits/stdc++.h>

using namespace std;

typedef long long ll;
typedef pair < int , int > pii;
typedef pair < ll , ll > pll;

#define X first
#define Y second
#define mp make_pair
#define pb push_back
#define all(x) (x).begin(), (x).end()
#define Sz(x) (int)(x.size())

const int N = 5e5 + 10;
const int L = 20;

int h[N], par[N][L], sum[N], P[N], sz[N], deg[N];
vector < int > adj[N], vec[N];

void DFS(int v, int p){
    par[v][0] = p;
    for (int i = 1; i < L; i ++)
        par[v][i] = par[par[v][i - 1]][i - 1];
    for (int u : adj[v])
        if (u != p)
            h[u] = h[v] + 1, DFS(u, v);
}

void calDFS(int v, int p){
    for (int u : adj[v])
        if (u != p)
            calDFS(u, v), sum[v] += sum[u];
    sum[v] ++;
}

int getPar(int v, int t){
    for (int i = 0; i < L; i ++)
        if (t >> i & 1)
            v = par[v][i];
    return v;
}

int getLca(int u, int v){
    if (h[v] < h[u])
        swap(u, v);
    v = getPar(v, h[v] - h[u]);
    if (u == v)
        return v;
    for (int i = L - 1; i >= 0; i --)
        if (par[u][i] != par[v][i])
            u = par[u][i], v = par[v][i];
    return par[v][0];
}

int get(int v){
    return (P[v] == v ? v : P[v] = get(P[v]));
}

void merge(int u, int v){
    u = get(u);
    v = get(v);
    if (u == v)
        return ;
    if (sz[v] < sz[u])
        swap(u, v);
    P[u] = v;
    sz[v] += sz[u];
}

int main(){
    ios::sync_with_stdio(0);
    cin.tie(0); cout.tie(0);
    int n, k;
    cin >> n >> k;
    for (int i = 1; i < n; i ++){
        int u, v;
        cin >> u >> v;
        adj[u].pb(v);
        adj[v].pb(u);
    }
    for (int i = 1; i <= n; i ++){
        int c; cin >> c;
        vec[c].pb(i);
    }
    DFS(1, 0);
    for (int i = 1; i <= k; i ++){
        int lca = vec[i][0];
        for (int j = 1; j < Sz(vec[i]); j ++)
            lca = getLca(lca, vec[i][j]);
        sum[lca] -= Sz(vec[i]);
    }
    calDFS(1, 0);
    for (int i = 0; i < N; i ++)
        P[i] = i, sz[i] = 1;
    for (int i = 2; i <= n; i ++)
        if (sum[i] != 0)
            merge(i, par[i][0]);
    for (int i = 2; i <= n; i ++)
        if (sum[i] == 0)
            deg[get(i)] ++, deg[get(par[i][0])] ++;
    int cnt = 0;
    for (int i = 1; i <= n; i ++)
        if (deg[i] == 1)
            cnt ++;
    cout << (cnt + 1) / 2;
    return 0;
}
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...