Submission #1092159

#	Time	Username	Problem	Language	Result	Execution time	Memory
1092159		vjudge1	Capital City (JOI20_capital_city)	C++17	41 / 100	3076 ms	32692 KiB

capital_city

#include <bits/stdc++.h>
using namespace std;
#define fi first
#define se second
#pragma GCC optimize("Ofast,unroll-loops")
//#pragma GCC target("avx2,popcnt,lzcnt,abm,bmi,bmi2,fma,tune=native")
#define all(x) x.begin(),x.end()
#define pb push_back
#define for1(i,x,n) for(int i=x;i<=n;i++)
#define for2(i,x,n) for(int i=x;i>=n;i--)
//#define int ll

typedef long long ll;
typedef pair<int,int> pii;

const ll maxn=2e5+69;
const ll mod=1e9+7;
const ll inf=1e18;

int n,k,c[maxn];
vector<int> adj[maxn],L[maxn];

int sz[maxn],par[maxn];
bool flag[maxn],vis[maxn],used[maxn],dd[maxn];

void dfs_sz(int u,int pre)
{
    sz[u]=1;
    for(int v: adj[u])
    {
        if (v==pre||used[v]) continue;
        dfs_sz(v,u);
        sz[u]+=sz[v];
    }
}
int Find_cen(int u,int pre,int ts)
{
    for(int v:adj[u])
    {
        if (v==pre||used[v]) continue;
        if (sz[v]*2>ts) return Find_cen(v,u,ts);
    }
    return u;
}
void dfs_on(int u,int pre)
{
    flag[u]=1;
    par[u]=pre;
    vis[u]=0;
    for(int v:adj[u])
    {
        if (v==pre||used[v]) continue;
        dfs_on(v,u);
    }
}
void dfs_off(int u,int pre)
{
    flag[u]=0;
    vis[u]=0;
    dd[c[u]]=0;
    for(int v:adj[u])
    {
        if (v==pre||used[v]) continue;
        dfs_off(v,u);
    }
}
vector<int>Q;
int centroid(int u)
{
    dfs_sz(u,-1);
    u=Find_cen(u,-1,sz[u]);
    used[u]=1;
    dfs_on(u,-1);
    Q.clear();
    for(int v:L[c[u]]) Q.pb(v);
    dd[c[u]]=1;
    int res=0;
    while (!Q.empty())
    {
        int u=Q.back();
        Q.pop_back();
        if (vis[u]) continue;
        if (!flag[u])
        {
            res=n;
            break;
        }
        vis[u]=1;
        if (par[u]==-1) continue;
        if (!dd[c[par[u]]])
        {
            for(int v:L[c[par[u]]])
            {
                if (!flag[v])
                {
                    res=n;
                    break;
                }
                Q.pb(v);
            }
            if (!Q.empty() && !flag[Q.back()]) break;
            dd[c[par[u]]]=1;
            res++;
        }
    }
//    Q.clear();
    dfs_off(u,-1);
    for(int v: adj[u]) if (!used[v]) res=min(res,centroid(v));
    return res;
}
void sol()
{
    cin >> n>>k;
    for1(i,1,n-1)
    {
        int u,v; cin >> u>>v;
        adj[u].pb(v);
        adj[v].pb(u);
    }
    for1(i,1,n)
    {
        cin >>c[i];
        L[c[i]].pb(i);
    }
    cout << centroid(1);
}

int32_t main()
{
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);
//    freopen("wall.inp","r",stdin);
//    freopen("wall.out","w",stdout);
    int t=1;//cin >> t;
    while (t--)
    {
//        cerr<<"avdsava "<< t<<'\n';
        sol();
    }
}
/*
1
3
3 2 3 1 1 2
4 1
6 1
3 5
5 1
2 4
*/

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