Submission #1092166

#TimeUsernameProblemLanguageResultExecution timeMemory
1092166vjudge1Capital City (JOI20_capital_city)C++17
100 / 100
313 ms41572 KiB


#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+13;
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],cnt[maxn];
bool used[maxn],dd[maxn];

void dfs_sz(int u,int pre)
{
    sz[u]=1;
    cnt[c[u]]=0;
    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)
{
    par[u]=pre;
    dd[c[u]]=0;
    cnt[c[u]]++;
    for(int v:adj[u])
    {
        if (v==pre||used[v]) continue;
        dfs_on(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();
    Q.pb(c[u]);
    dd[c[u]]=1;
    int res=0;
    while (!Q.empty())
    {
        int cu=Q.back();
        Q.pop_back();
        if (cnt[cu]!=L[cu].size())
        {
            res=n;
            break;
        }
        for(int v:L[cu])
        {
            if (par[v]==-1) continue;
            if (!dd[c[par[v]]])
            {
                Q.pb(c[par[v]]);
                dd[c[par[v]]]=1;
                res++;
            }
        }
    }
    Q.clear();
//    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
*/

Compilation message (stderr)

capital_city.cpp: In function 'int centroid(int)':
capital_city.cpp:75:20: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   75 |         if (cnt[cu]!=L[cu].size())
      |             ~~~~~~~^~~~~~~~~~~~~~
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...