제출 #991895

#	제출 시각	아이디	문제	언어	결과	실행 시간	메모리
991895		MarwenElarbi	Mergers (JOI19_mergers)	C++17	100 / 100	1335 ms	260016 KiB

mergers

#include<bits/stdc++.h>
using namespace std;
#pragma GCC optimize("O3")
#pragma GCC optimize("unroll-loops")
#define fi first
#define se second
#define ll long long
#define pb push_back
const int nax=5e5+5;
vector<int> adj[nax];
set<int> stl[nax];
vector<int> sz(nax);
vector<int> sum(nax);
vector<int> s(nax);
vector<int> colours(nax);
int ans=0;
vector<int> parent(nax);
vector<int> adj2[nax];
vector<bool> vis(nax);
int find(int x){
    if(parent[x]==x) return x;
    return parent[x]=find(parent[x]);
}
bool sameset(int x,int y){
    return find(x)==find(y);
}
void joinset(int x,int y){
    x=find(x);
    y=find(y);
    parent[x]=y;
}
void precompute(int x,int p){
    stl[x].insert(s[x]);
    sum[x]=colours[s[x]];
    sz[x]=1;
    for(auto u:adj[x]){
        if(u==p) continue;
        precompute(u,x);
        sz[x]+=sz[u];
        if(stl[u].size()>stl[x].size()){
            sum[x]=sum[u];
            swap(stl[u],stl[x]);
        }    
        for(auto i:stl[u]){
            if(stl[x].count(i)) continue;
            stl[x].insert(i);
            sum[x]+=colours[i];
        }
    }
    if(p!=-1&&sum[x]>sz[x]){
        joinset(x,p);
    }
}
void nab(int x,int p){
    for(auto u:adj[x]){
        if(u==p) continue;
        if(!sameset(x,u)){
            int curx=find(x);
            int cury=find(u);
            adj2[curx].pb(cury);
            adj2[cury].pb(curx);
        }
        nab(u,x);
    }
}
void dfs(int x,int p){
    //cout <<x<<" "<<adj2[x].size()<<endl;
    if(adj2[x].size()==1) ans++;
    for(auto u:adj2[x]){
        if(u==p) continue;
        dfs(u,x);
    }
    return;
}
int main() {
    int n,k;
    cin>>n>>k;
    for (int i = 0; i < n-1; ++i)
    {
        int a,b;
        cin>>a>>b;
        a--;b--;
        adj[a].pb(b);
        adj[b].pb(a);
    }
    for (int i = 0; i < n; ++i)
    {
        parent[i]=i;
    }
    for (int i = 0; i < n; ++i)
    {
        cin>>s[i];
        s[i]--;
        colours[s[i]]++;
    }
    precompute(0,-1);
    nab(0,-1);
    for (int i = 0; i < n; ++i)
    {
        if(adj2[i].size()>0){
            dfs(i,-1);
            break;
        }
    }
    cout <<(ans+1)/2<<endl;
}

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