Submission #223921

#	Time	Username	Problem	Language	Result	Execution time	Memory
223921		medk	Capital City (JOI20_capital_city)	C++14	100 / 100	544 ms	84536 KiB

capital_city

#include <bits/stdc++.h>

#define ld long double
#define ll long long
#define pb push_back
#define x first
#define y second
#define all(x) x.begin(),x.end()
#define sz(x) (int)(x.size())

using namespace std;

const int MAXN=200000;

int n,k;
int lca[20][500000];
vector<vector<int>> g(MAXN);
vector<set<int>> unvis(MAXN);
vector<int> eul,lvl(MAXN),in(MAXN),par(MAXN),color,colorLCA,dsu,dsu2,sz,highestV,doneColors(MAXN);
vector<int> vis(MAXN);

int getp(int x)
{
    if(x==dsu[x]) return x;
    return dsu[x]=getp(dsu[x]);
}

void connect(int u, int v)
{
    u=getp(u), v=getp(v);
    if(u==v) return;
    if(sz[v]>sz[u]) swap(u,v);
    dsu[v]=u;
    sz[u]+=sz[v];
    doneColors[u]+=doneColors[v];
}

int getp2(int x)
{
    if(x==dsu2[x]) return x;
    return dsu2[x]=getp2(dsu2[x]);
}
void connect2(int u, int v)
{
    u=getp2(u), v=getp2(v);
    if(u==v) return;
    if(rand()&1) swap(u,v);
    dsu2[v]=u;
    if(lvl[highestV[u]]>lvl[highestV[v]]) highestV[u]=highestV[v];
}
void makeeuler(int u, int pr, int d)
{
    par[u]=pr;
    in[u]=sz(eul);
    lvl[u]=d;
    eul.pb(u);
    for(int v:g[u])
    {
        if(v==pr) continue;
        makeeuler(v,u,d+1);
        eul.pb(u);
    }
}

int getlca(int u, int v)
{
    int l=in[u], r=in[v];
    if(l>r) swap(l,r);
    int lg=log2(r-l+1);
    return (lvl[lca[lg][l]]<=lvl[lca[lg][r-(1<<lg)+1]]?lca[lg][l]:lca[lg][r-(1<<lg)+1]);
}

int main()
{
    //freopen("test.in","r",stdin);
    ios::sync_with_stdio(0); cin.tie(0);
    cin>>n>>k;
    for(int i=0;i<k;i++) dsu.pb(i), sz.pb(1);
    for(int i=0;i<n-1;i++)
    {
        int u,v; cin>>u>>v;
        u--;v--;
        g[u].pb(v);
        g[v].pb(u);
    }
    makeeuler(0,-1,0);
    for(int i=0;i<sz(eul);i++) lca[0][i]=eul[i];
    for(int i=1;i<20;i++)
        for(int j=0;j+(1<<i)<=sz(eul);j++)
            lca[i][j] = (lvl[lca[i-1][j]]<=lvl[lca[i-1][j+(1<<(i-1))]]?lca[i-1][j]:lca[i-1][j+(1<<(i-1))]);
    for(int i=0;i<n;i++)
    {
        dsu2.pb(i), highestV.pb(i);
        int c; cin>>c; c--;
        color.pb(c);
        unvis[c].insert(i);
    }
    int ans=1e9;
    set<pair<int,int>> q;
    for(int i=0;i<k;i++)
    {
        int lc=*unvis[i].begin();
        for(int v:unvis[i])
            lc=getlca(lc,v);
        colorLCA.pb(lc);
        q.insert({-lvl[lc],lc});
    }
    for(auto p:q)
    {
        int candidLCA=p.y;
        if(candidLCA!=colorLCA[color[candidLCA]]) continue;
        queue<int> qu;
        for(int v:unvis[color[candidLCA]])
        {
            qu.push(v);
            if(v!=candidLCA) vis[v]=1;
        }
        unvis[color[candidLCA]].clear();
        doneColors[getp(color[candidLCA])]++;
        while(!qu.empty())
        {
            int v=qu.front(); qu.pop();
            while(v!=candidLCA)
            {
                if(vis[v]==2) break;
                connect(color[candidLCA],color[v]);
                if(lvl[colorLCA[color[v]]]>=lvl[candidLCA])
                {
                    int bef=sz(unvis[color[v]]);
                    unvis[color[v]].erase(v);
                    if(sz(unvis[color[v]])==0 && bef==1) doneColors[getp(color[v])]++;
                    vis[v]=2;
                    for(int t:unvis[color[v]])
                    {
                        if(vis[t]) continue;
                        vis[t]=1;
                        qu.push(t);
                        unvis[color[t]].erase(t);
                        if(sz(unvis[color[t]])==0) doneColors[getp(color[t])]++;
                    }
                }
                else break;
                connect2(v,par[v]);
                v=highestV[getp2(par[v])];
                if(vis[v]==1) break;
            }
        }
        int rep=getp(color[candidLCA]);
        if(doneColors[rep]==sz[rep]) ans=min(ans,sz[rep]-1);
    }
    cout<<ans<<'\n';
    return 0;
}

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