Submission #966221

#	Time	Username	Problem	Language	Result	Execution time	Memory
966221		12345678	Capital City (JOI20_capital_city)	C++17	100 / 100	310 ms	68000 KiB

capital_city

#include <bits/stdc++.h>

using namespace std;

const int nx=2e5+5, kx=18;

int n, k, u, v, t[nx], lvl[nx], pa[nx][kx], h[nx], vs[nx], rt[nx], sz[nx], ans=INT_MAX;
vector<int> d[nx], c[nx], rv[nx], g[nx];
stack<int> s;
vector<pair<int, int>> edg;

void dfs(int u, int p)
{
    lvl[u]=lvl[p]+1;
    pa[u][0]=p;
    for (int i=1; i<kx; i++) pa[u][i]=pa[pa[u][i-1]][i-1];
    for (auto v:d[u]) if (v!=p) dfs(v, u);
}

int lca(int u, int v)
{
    if (lvl[u]>lvl[v]) swap(u, v);
    for (int i=kx-1; i>=0; i--) if (lvl[pa[v][i]]>=lvl[u]) v=pa[v][i];
    if (u==v) return u;
    for (int i=kx-1; i>=0; i--) if (pa[u][i]!=pa[v][i]) u=pa[u][i], v=pa[v][i];
    return pa[u][0];
}

void dfs1(int u, int p)
{
    vs[u]=1;
    for (auto v:c[u]) if (v!=p&&!vs[v]) dfs1(v, u);
    s.push(u);
}

void dfs2(int u, int p, int hd)
{
    vs[u]=1;
    rt[u]=hd;
    sz[hd]++;
    for (auto v:rv[u]) if (v!=p&&!vs[v]) dfs2(v, u, hd);
}

int main()
{
    cin.tie(NULL)->sync_with_stdio(false);
    cin>>n>>k;
    for (int i=1; i<n; i++) cin>>u>>v, d[u].push_back(v), d[v].push_back(u);
    for (int i=1; i<=n; i++) cin>>t[i], g[t[i]].push_back(i);
    dfs(1, 1);
    for (int i=1; i<=k; i++)
    {
        int hd=g[i][0];
        for (auto u:g[i]) hd=lca(hd, u);
        h[i]=hd;
    }
    for (int i=1; i<=n; i++) if (i!=h[t[i]]&&t[pa[i][0]]!=t[i]) c[t[i]].push_back(t[pa[i][0]]), rv[t[pa[i][0]]].push_back(t[i]), edg.push_back({t[i], t[pa[i][0]]}); //cout<<"edge "<<t[i]<<' '<<t[pa[i][0]]<<'\n';
    for (int i=1; i<=k; i++) if (!vs[i]) dfs1(i, i);
    for (int i=1; i<=k; i++) vs[i]=0;
    while (!s.empty())
    {
        u=s.top();
        //cout<<"stack "<<u<<'\n';
        s.pop();
        if (vs[u]) continue;
        dfs2(u, u, u);
    }
    //for (int i=1; i<=n; i++) cout<<"rt "<<i<<' '<<rt[i]<<'\n';
    for (auto [u, v]:edg) if (rt[u]!=rt[v]) sz[u]=1e9;
    for (int i=1; i<=k; i++) ans=min(ans, sz[rt[i]]);
    cout<<ans-1;
}

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