Submission #1088820

#	Time	Username	Problem	Language	Result	Execution time	Memory
1088820		phong	Capital City (JOI20_capital_city)	C++17	100 / 100	421 ms	42868 KiB

capital_city

#include<bits/stdc++.h>
#define ll long long
const int nmax = 2e5 + 5, N = 1e6;
const ll oo = 1e9 + 5, base = 311;
const int lg = 62, tx = 26;
const ll mod = 1e9 + 7, mod2 = 1e9+3;
#define pii pair<ll, ll>
#define fi first
#define se second
#define endl "\n"
#define debug(a, n) for(int i = 1; i <= n; ++i) cout << a[i] << ' '; cout << "\n";
#define NAME code
using namespace std;


int n, k;
vector<int> adj[nmax];
vector<int> mau[nmax];
int idx[nmax], c[nmax];
int sz[nmax], vis[nmax], cntw;

void dfs_1(int u, int p){
    sz[u] = 1;
    for(auto v : adj[u]){
        if(vis[v] || v == p) continue;
        dfs_1(v, u);
        sz[u] += sz[v];
    }
}

int Find(int u, int p){
    for(auto v : adj[u]){
        if(vis[v] || v == p) continue;
        if(sz[v] > cntw / 2) return Find(v, u);
    }
    return u;
}
int rev[nmax], ok[nmax], par[nmax];

int res = 0;
void Init(int u, int p, int val){
    ok[u] = val;
    for(auto v : adj[u]){
        if(vis[v] || v == p) continue;
        par[v] = u;
        Init(v, u, val);
    }
}
struct holder{
    int vis[nmax];
    vector<int> tmp;
    void add(int u){
        vis[u] = 1;
        tmp.push_back(u);
    }
    void clear(){
        for(auto p : tmp) vis[p] = 0;
        tmp.clear();
    }
}one;
queue<int> q;
void solve(int u){
    int root = u;
    while(q.size()) q.pop();
    for(auto p : mau[c[root]]){
        if(!ok[p]){
            res = oo;
            return;
        }
        q.push(p);
        one.add(p);
    }
    while(q.size()){
        int u = q.front();q.pop();
        if(u == root) continue;
        int v = par[u];
//        cout << u << ' ' << v << endl;
        if(one.vis[v]) continue;
        res++;
        for(auto p : mau[c[v]]){
            if(!ok[p]){
                res = oo;
                return;
            }
            q.push(p);
            one.add(p);
        }
     }
}
int centroid(int u){
    dfs_1(u, -1); cntw = sz[u];
    int root = Find(u,-1);vis[root] = 1;
    res = 0;
    one.clear();
    Init(root, -1, 1);
    solve(root);
    Init(root, -1, 0);
    int cur = res;
    for(auto v : adj[root]){
        if(!vis[v])cur = min(cur, centroid(v));
    }
    return cur;
}
main(){
    ios_base::sync_with_stdio(0);
    cin.tie(0);cout.tie(0);
//    freopen("code.inp", "r", stdin);
//    freopen("code.out", "w", stdout);
    cin >> n >> k;
    for(int i = 1; i < n; ++i){
        int u, v;
        cin >> u >> v;
        adj[u].push_back(v);
        adj[v].push_back(u);
    }
    for(int i = 1; i <= n; ++i){
        cin >> c[i];
        mau[c[i]].push_back(i);
    }
    cout << centroid(1);
}
/*
7 3
3 2
1 3
7 3
5 2
6 2
4 7
2 1 3 1 2 3 2

*/

Compilation message (stderr)

capital_city.cpp:104:1: warning: ISO C++ forbids declaration of 'main' with no type [-Wreturn-type]
  104 | main(){
      | ^~~~

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