Submission #1279260

#	Time	Username	Problem	Language	Result	Execution time	Memory
1279260		denislav	Capital City (JOI20_capital_city)	C++20	100 / 100	847 ms	36284 KiB

capital_city

# include <iostream>
# include <vector>
# include <queue>
using namespace std;

const int MAX=2e5+11,INF=1e9;

int n,k;
vector<int> adj[MAX];
vector<int> colors[MAX];
int col[MAX];

bool killed[MAX];
int sz[MAX];

void dfs_sz(int curr, int par)
{
    sz[curr]=1;
    for(int nxt: adj[curr])
    {
        if(nxt==par or killed[nxt]) continue;
        dfs_sz(nxt,curr);
        sz[curr]+=sz[nxt];
    }
}

int get_centroid(int curr, int par, int N)
{
    for(int nxt: adj[curr])
    {
        if(nxt==par or killed[nxt]) continue;
        if(sz[nxt]*2>N) return get_centroid(nxt,curr,N);
    }
    return curr;
}

bool state[MAX];

void dfs_state(int curr, int par)
{
    state[curr]^=1;
    for(int nxt: adj[curr])
    {
        if(nxt==par or killed[nxt]) continue;
        dfs_state(nxt,curr);
    }
}

bool vis[MAX];
bool vis_col[MAX];
int h[MAX];
int up[MAX];

void dfs_init(int curr, int par)
{
    for(int nxt: adj[curr])
    {
        if(nxt==par or killed[nxt]) continue;
        h[nxt]=h[curr]+1;
        up[nxt]=curr;
        dfs_init(nxt,curr);
    }
}

int solve(int start)
{
    h[start]=0;
    dfs_init(start,0);

    queue<int> q;
    for(int x: colors[col[start]])
    {
        if(!state[x])
        {
            vis[start]=0;
            vis_col[col[start]]=0;
            return INF;
        }
        if(x!=start) q.push(x);
        else
        {
            vis[start]=1;
            vis_col[col[start]]=1;
        }
    }

    bool f=1;
    int ans=1;
    vector<int> toRem;
    while(q.size()>0)
    {
        int curr=q.front();
        q.pop();

        while(!vis[curr])
        {
            toRem.push_back(curr);
            vis[curr]=1;
            if(vis_col[col[curr]]==0)
            {
                ans++;
                vis_col[col[curr]]=1;
                for(int nxt: colors[col[curr]])
                {
                    if(!state[nxt])
                    {
                        f=0;
                        break;
                    }

                    q.push(nxt);
                }
            }
            curr=up[curr];
        }

        if(!f) break;
    }

    toRem.push_back(start);
    for(int x: toRem)
    {
        vis[x]=0;
        vis_col[col[x]]=0;
    }

    if(f) return ans;
    else return INF;
}

int ANS=INF;

void cd(int curr)
{
    dfs_sz(curr,0);
    curr=get_centroid(curr,0,sz[curr]);

    dfs_state(curr,0);
    ANS=min(ANS,solve(curr));
    dfs_state(curr,0);

    killed[curr]=1;
    for(int nxt: adj[curr])
    {
        if(!killed[nxt]) cd(nxt);
    }
}

int main()
{
    ios_base::sync_with_stdio(0);
    cin.tie(0);cout.tie(0);

    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>>col[i];
        colors[col[i]].push_back(i);
    }

    cd(1);

    cout<<ANS-1<<"\n";

    return 0;
}

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