Submission #1014315

#	Time	Username	Problem	Language	Result	Execution time	Memory
1014315		alexdd	Capital City (JOI20_capital_city)	C++17	100 / 100	419 ms	42620 KiB

capital_city

#include<bits/stdc++.h>
using namespace std;
const int INF = 1e9;
int n,k,rez=INF;
vector<int> con[200005];
int color[200005];
int fr[200005];
int siz[200005];
bool iss[200005];
bool hascol[200005];
bool bagat[200005];
vector<int> ofcol[200005];
vector<int> active;
int parent[200005];
void get_sizes(int nod, int par)
{
    if(ofcol[color[nod]].empty()) active.push_back(color[nod]);
    ofcol[color[nod]].push_back(nod);
    siz[nod]=1;
    bagat[color[nod]]=0;
    hascol[nod]=0;
    for(auto adj:con[nod])
    {
        if(adj!=par && !iss[adj])
        {
            get_sizes(adj,nod);
            siz[nod]+=siz[adj];
        }
    }
}
int get_centroid(int nod, int par, int tot)
{
    for(auto adj:con[nod])
        if(adj!=par && !iss[adj] && 2*siz[adj]>tot)
            return get_centroid(adj,nod,tot);
    return nod;
}
void dfs_parent(int nod)
{
    for(auto adj:con[nod])
    {
        if(!iss[adj] && adj!=parent[nod])
        {
            parent[adj]=nod;
            dfs_parent(adj);
        }
    }
}
int cate;
void baga_color(int col);
void baga_nod(int nod)
{
    if(hascol[nod])
        return;
    hascol[nod]=1;
    baga_color(color[nod]);
    if(parent[nod]!=-1 && !hascol[parent[nod]])
        baga_nod(parent[nod]);
}
void baga_color(int col)
{
    if(bagat[col])
        return;
    cate++;
    if((int)ofcol[col].size() != fr[col])
        cate=INF;
    bagat[col]=1;
    for(auto x:ofcol[col])
    {
        baga_nod(x);
    }
}
void centroid_decomposition(int c)
{
    get_sizes(c,-1);
    c = get_centroid(c,-1,siz[c]);

    if((int)ofcol[color[c]].size()==fr[color[c]])
    {
        parent[c]=-1;
        dfs_parent(c);
        cate=0;
        baga_color(color[c]);
        rez = min(rez, cate-1);
    }

    for(auto x:active)
        ofcol[x].clear();
    active.clear();

    iss[c]=1;
    for(auto adj:con[c])
        if(!iss[adj])
            centroid_decomposition(adj);
}
signed main()
{
    cin>>n>>k;
    int a,b;
    for(int i=1;i<n;i++)
    {
        cin>>a>>b;
        con[a].push_back(b);
        con[b].push_back(a);
    }
    for(int i=1;i<=n;i++)
    {
        cin>>color[i];
        fr[color[i]]++;
    }
    centroid_decomposition(1);
    cout<<rez;
    return 0;
}

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