Submission #1095091

#	Time	Username	Problem	Language	Result	Execution time	Memory
1095091		pokmui9909	Mergers (JOI19_mergers)	C++17	70 / 100	3080 ms	214120 KiB

mergers

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
 
#define x first
#define y second
 
ll N, K, A[500005], S[500005], vis[500005];
vector<ll> T[500005];
 
ll Size(ll u, ll p){
    S[u] = 1;
    for(auto v : T[u]){
        if(v != p && !vis[v]) S[u] += Size(v, u);
    }
    return S[u];
}
ll Cent(ll u, ll p, ll n){
    for(auto v : T[u]){
        if(v != p && !vis[v] && S[v] > n / 2) return Cent(v, u, n);
    }
    return u;
}   
 
ll Cnt[500005];
vector<ll> G[500005], H[500005];
 
void dfs(ll u, ll p, ll d){
    Cnt[A[u]] += d;
    for(auto v : T[u]){
        if(v == p || vis[v]) continue;
        dfs(v, u, d);
    }
}
ll Comp(ll u, ll p, ll c){
    ll ok = (Cnt[A[u]] > 0);
    for(auto v : T[u]){
        if(v == p || vis[v]) continue;
        if(Comp(v, u, c)) ok = 1;
    }
    if(ok){
        G[c].push_back(A[u]); G[A[u]].push_back(c);
    }
    return ok;
}
void dnc(ll _u){
    ll u = Cent(_u, -1, Size(_u, -1));
    vis[u] = 1;
 
    Cnt[A[u]]++;
    for(auto v : T[u]){
        if(!vis[v]) dfs(v, u, 1);
    }
    for(auto v : T[u]){
        if(vis[v]) continue;
        dfs(v, u, -1);
        Comp(v, u, A[u]);
        dfs(v, u, 1);
    }
    for(auto v : T[u]){
        if(!vis[v]) dfs(v, u, -1);
    }
    Cnt[A[u]]--;
 
    for(auto v : T[u]){
        if(!vis[v]) dnc(v);
    }
}
 
ll deg[500005], Num[500005];
 
void f(ll u){
    for(auto v : G[u]){
        if(Num[v]) continue;
        Num[v] = Num[u]; f(v);
    }
}
 
ll chk[500005];
 
int main(){
    cin.tie(0) -> sync_with_stdio(0);
 
    cin >> N >> K;
    for(ll i = 1; i < N; i++){
        ll u, v; cin >> u >> v;
        T[u].push_back(v); T[v].push_back(u);
    }
    for(ll i = 1; i <= N; i++){
        cin >> A[i];
    }
    dnc(1);
    for(ll i = 1; i <= N; i++){
        for(auto j : G[i]){
            if(chk[j]) continue;
            H[i].push_back(j);
            chk[j] = 1;
        }
        for(auto j : H[i]) chk[j] = 0;
    }
    for(ll i = 1, j = 0; i <= N; i++){
        if(Num[i]) continue;
        Num[i] = ++j; f(i);
    }
    for(ll i = 1; i <= N; i++){
        for(auto j : T[i]){
            if(i > j) continue;
            ll a = Num[A[i]], b = Num[A[j]];
            if(a != b){
                deg[a]++, deg[b]++;
            }
        }
    }
    ll f = 0;
    for(ll i = 1; i <= N; i++){
        if(deg[i] == 1) f++;
    }
    cout << (f + 1) / 2;
}

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