Submission #359153

#	Time	Username	Problem	Language	Result	Execution time	Memory
359153		Nima_Naderi	Mergers (JOI19_mergers)	C++14	100 / 100	1402 ms	128096 KiB

mergers

///In the name of GOD
//#pragma GCC optimize("O2")
#include<bits/stdc++.h>
using namespace std;

//typedef long long ll;
typedef int ll;
const ll MXN = 5e5 + 10;
ll n, k, leaf;
ll A[MXN], has[MXN], dad[MXN], hd[MXN];
ll Par[MXN], Sz[MXN], dis[MXN];
ll sub[MXN], hvs[MXN];
vector<ll> adj[MXN], G[MXN], Col[MXN];
vector<pair<ll, ll>> E;
ll Find(ll x){
    return (x == Par[x] ? x : Par[x] = Find(Par[x]));
}
void Union(ll x, ll y){
    x = Find(x), y = Find(y);
    if(x == y) return;
    if(Sz[x] < Sz[y]) swap(x, y);
    Sz[x] += Sz[y], Par[y] = x;
}
void prep(ll u, ll par){
    sub[u] ++, dad[u] = par;
    if(A[u]) Col[A[u]].push_back(u);
    for(auto v : adj[u]){
        if(v == par) continue;
        dis[v] = dis[u] + 1;
        prep(v, u);
        sub[u] += sub[v];
        if(sub[v] > sub[hvs[u]]) hvs[u] = v;
    }
}
void dfs(ll u, ll par, ll head){
    hd[u] = head;
    if(hvs[u]) dfs(hvs[u], u, head);
    for(auto v : adj[u]){
        if(v == par || v == hvs[u]) continue;
        dfs(v, u, v);
    }
}
void flood(ll u, ll par){
    for(auto v : adj[u]){
        if(v == par) continue;
        flood(v, u), has[u] += has[v];
        if(has[v]) Union(u, v);
    }
}
ll LCA(ll u, ll v){
    while(hd[u] != hd[v]){
        if(dis[hd[u]] > dis[hd[v]]) swap(u, v);
        v = dad[hd[v]];
    }
    if(dis[u] > dis[v]) swap(u, v);
    return u;
}
int main(){
    ios::sync_with_stdio(0);cin.tie(0); cout.tie(0);
    cin >> n >> k;
    iota(Par + 1, Par + n + 1, 1), fill(Sz + 1, Sz + n + 1, 1);
    for(int i = 1; i < n; i ++){
        ll u, v; cin >> u >> v;
        adj[u].push_back(v), adj[v].push_back(u);
        E.push_back({u, v});
    }
    for(int i = 1; i <= n; i ++) cin >> A[i];
    prep(1, 0), dfs(1, 0, 1);
    for(int i = 1; i <= k; i ++){
        for(int j = 1; j < Col[i].size(); j ++){
            ll u = Col[i][j - 1], v = Col[i][j];
            has[u] ++, has[v] ++, has[LCA(u, v)] -= 2;
        }
    }
    flood(1, 0);
    for(auto e : E){
        ll u, v; tie(u, v) = e;
        u = Find(u), v = Find(v);
        if(u == v) continue;
        G[u].push_back(v), G[v].push_back(u);
    }
    for(int i = 1; i <= n; i ++) leaf += (Find(i) == i && G[i].size() == 1);
    cout << (leaf + 1) / 2 << '\n';
    return 0;
}
/*!
    HE'S AN INSTIGATOR,
    ENEMY ELIMINATOR,
    AND WHEN HE KNOCKS YOU BETTER
    YOU BETTER LET HIM IN.
*/
//! N.N

Compilation message (stderr)

mergers.cpp: In function 'int main()':
mergers.cpp:70:26: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   70 |         for(int j = 1; j < Col[i].size(); j ++){
      |                        ~~^~~~~~~~~~~~~~~

• Subtask 1: 10 / 10
• Subtask 2: 24 / 24
• Subtask 3: 14 / 14
• Subtask 4: 22 / 22
• Subtask 5: 30 / 30