Submission #1182997

#TimeUsernameProblemLanguageResultExecution timeMemory
1182997clemmy14Mergers (JOI19_mergers)C++20
100 / 100
808 ms166196 KiB
#include<bits/stdc++.h>
using namespace std;

struct DSU {
    vector<int> e; 
    DSU(int N) {e = vector<int>(N, -1);}
    int get(int x) {return e[x] < 0 ? x : e[x]=get(e[x]);}
    int size(int x) {return -e[get(x)];}
    bool unite(int a, int b) {
        a=get(a); b=get(b);
        if(a == b) return false;
        if(e[a] > e[b]) swap(a, b);
        e[a]+=e[b]; e[b]=a;
        return true;
    }
};

int timer=0;
vector<int> depth, timeIn;
vector<vector<int>> adj, parent;

void dfs1(int node, int prev, int dis) {
    timeIn[node]=timer++;
    depth[node]=dis;
    parent[node][0]=prev;
    for(auto child : adj[node]) if(child != prev) dfs1(child, node, dis+1);
}

int ancestor(int x, int k) {
    int cnt=0;
    while(k&&x) {
        if(k&1) x = parent[x][cnt];
        k=k>>1; cnt++;
    }
    return x;
}

int lca(int a, int b) {
    if(depth[a] < depth[b])swap(a, b);
    a=ancestor(a, depth[a]-depth[b]);
    if(a==b) return a;
    for(int k=19; k>=0; k--) {
        int aa=parent[a][k];
        int bb=parent[b][k];
        if(aa!=bb) {
            a=aa; b=bb;
        }
    }
    return parent[a][0];
}

vector<int> connections;

void dfs2(int node, int prev, DSU&dsu) {
    for(auto child : adj[node]) if(prev != child) {
        int u=dsu.get(node), v=dsu.get(child);
        if(u != v) {
            connections[u]++;
            connections[v]++;
        }
        dfs2(child, node, dsu);
    }
}

signed main() {
    int n, k; cin >> n >> k;
    adj = vector<vector<int>>(n+1);
    for(int i=1; i<n; i++) {
        int a, b; cin >> a >> b; 
        adj[a].push_back(b);
        adj[b].push_back(a);
    }

    DSU dsu(n+1);
    vector<int> type(n+1);
    vector<vector<pair<int, int>>> groups(k+1);
    timeIn = vector<int>(n+1);
    depth = vector<int>(n+1);
    parent = vector<vector<int>>(n+1, vector<int>(20, 0));

    dfs1(1, 0, 0);
    for(int k=1; k<20; k++) {
        for(int i=1; i<=n; i++) {
            if(parent[i][k-1]) parent[i][k] = parent[parent[i][k-1]][k-1];
        }
    }

    for(int i=1; i<=n; i++) {
        cin >> type[i];
        groups[type[i]].push_back({timeIn[i], i});
    }

    for(int i=1; i<=k; i++) {
        sort(groups[i].begin(), groups[i].end());
        for(int j=1; j<groups[i].size(); j++) {
            int curLca = lca(groups[i][j-1].second, groups[i][j].second);
            int cur=groups[i][j-1].second;
            while(cur != curLca) {
                dsu.unite(curLca, cur);
                cur=parent[cur][0];
            }
            cur=groups[i][j].second;
            while(cur != curLca) {
                dsu.unite(curLca, cur);
                cur=parent[cur][0];
            }
        }

    }
    connections=vector<int>(n+1, 0);
    dfs2(1, 0, dsu);
    int tt=0;
    for(auto x : connections) if(x == 1) tt++;
    cout << (tt+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...