Submission #1272233

#	Time	Username	Problem	Language	Result	Execution time	Memory
1272233		danglayloi1	Capital City (JOI20_capital_city)	C++20	100 / 100	586 ms	33244 KiB

capital_city

#include <bits/stdc++.h>
#define ii pair<int, int>
#define fi first
#define se second
#define inf 0x3f3f3f3f
using namespace std;
using ll = long long;
const ll mod=1e9+7;
const int nx=2e5+5;
int n, k, sz[nx], par[nx], id[nx], cnt[nx], ans=inf;
vector<int> adj[nx], ve, sam, col[nx];
bool del[nx], vis[nx];
int go(int p, int u)
{
    sz[u]=1;
    for(int v:adj[u])
        if(v!=p && !del[v]) sz[u]+=go(u, v);
    return sz[u];
}
int centroid(int p, int u, int tre)
{
    for(int v:adj[u])
        if(v!=p && !del[v] && sz[v]>tre/2) return centroid(u, v, tre);
    return u;
}
void dfs(int p, int u, int val)
{
    par[u]=p;
    cnt[id[u]]++;
    ve.emplace_back(u);
    if(id[u]==val) sam.emplace_back(u);
    for(int v:adj[u])
        if(v!=p && !del[v]) dfs(u, v, val);
}
void build(int u)
{
    int root=centroid(0, u, go(0, u));
    dfs(0, root, id[root]);
    if(cnt[id[root]]==(int)col[id[root]].size())
    {
        int dem=0;
        vis[id[root]]=1;
        bool ok=1;
        while(sam.size())
        {
            int x=par[sam.back()];
            sam.pop_back();
            if(x>0 && !vis[id[x]])
            {
                if(cnt[id[x]]<(int)col[id[x]].size())
                {
                    ok=0;
                    break;
                }
                dem++;
                vis[id[x]]=1;
                for(int y:col[id[x]])
                    sam.emplace_back(y);
            }
        }
        if(ok) ans=min(ans, dem);
    }
    for(int x:ve) vis[id[x]]=0, cnt[id[x]]=0;
    sam.clear();
    ve.clear();
    del[root]=1;
    for(int v:adj[root])
        if(!del[v]) build(v);
}
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].emplace_back(v);
        adj[v].emplace_back(u);
    }
    for(int i = 1; i <= n; i++)
        cin>>id[i], col[id[i]].emplace_back(i);
    build(1);
    cout<<ans;
}

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