Submission #314291

#	Time	Username	Problem	Language	Result	Execution time	Memory
314291		vipghn2003	Capital City (JOI20_capital_city)	C++14	100 / 100	1038 ms	36712 KiB

capital_city

#include<bits/stdc++.h>

using namespace std;

const int N=2e5+5;
int n,k,c[N],cnt[N],res,sz[N],par[N],num[N];
vector<int>adj[N],col[N],vec;
bool cenvis[N],colvis[N],vis[N];

void findsz(int u,int p=-1)
{
	sz[u]=1;
	for(auto&v:adj[u])
    {
		if (v!=p&&!cenvis[v])
		{
			findsz(v,u);
			sz[u]+=sz[v];
		}
	}
}

int find_cen(int u,int p,int m)
{
	for(auto&v:adj[u])
	{
		if(v!=p&&!cenvis[v])
        {
			if (sz[v]>m/2) return find_cen(v,u,m);
		}
	}
	return u;
}

void prep(int u,int p)
{
    vec.push_back(c[u]);
    num[c[u]]=0;
    par[u]=p;
    col[c[u]].push_back(u);
    vis[u]=false;
    colvis[c[u]]=false;
    for(auto&v:adj[u])
    {
        if(v!=p&&!cenvis[v]) prep(v,u);
    }
}

void centroid(int root=1,int p=-1)
{
    findsz(root,root);
	int u=find_cen(root,root,sz[root]);
	prep(u,u);
    queue<int>pq;
    pq.push(c[u]);
    colvis[c[u]]=true;
    int op=0;
    while(!pq.empty())
    {
        int color=pq.front();
        pq.pop();
        op++;
        for(auto&x:col[color])
        {
            while(!vis[x])
            {
                if(!colvis[c[x]]) pq.push(c[x]);
                vis[x]=true;
                colvis[c[x]]=true;
                num[c[x]]++;
                x=par[x];
            }
        }
    }
    bool kt=true;
    while(!vec.empty())
    {
        int color=vec.back();
        col[color].clear();
        vec.pop_back();
        if(colvis[color]&&num[color]!=cnt[color]) kt=false;
    }
    if(kt) res=min(res,op-1);
	cenvis[u]=true;
	for(auto&v:adj[u]) if(!cenvis[v]) centroid(v,u);
}

int main()
{
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);
    cin>>n>>k;
    for(int i=1;i<n;i++)
    {
        int u,v;
        cin>>u>>v;
        adj[u].push_back(v);
        adj[v].push_back(u);
    }
    for(int i=1;i<=n;i++)
    {
        cin>>c[i];
        cnt[c[i]]++;
    }
    res=1e9;
    centroid(1);
    cout<<res;
}

• Subtask 1: 1 / 1
• Subtask 2: 10 / 10
• Subtask 3: 30 / 30
• Subtask 4: 59 / 59