Submission #1088810

#	Time	Username	Problem	Language	Result	Execution time	Memory
1088810		phong	Capital City (JOI20_capital_city)	C++17	0 / 100	357 ms	92340 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];
int h[nmax], up[nmax][lg + 1];
void dfs(int u, int p){
    for(auto v : adj[u]){
        if(v == p) continue;
        h[v] = h[u] + 1;
        up[v][0] = u;
        for(int j = 1; j <= lg; ++j) up[v][j] = up[up[v][j - 1]][j - 1];
        dfs(v, u);
    }
}

int lca(int u, int v){
    if(h[u] != h[v]){
        if(h[u] < h[v]) swap(u, v);
        int k = h[u] - h[v];
        for(int j = 0; j <= __lg(k); ++j){
            if(k >> j & 1){
                u = up[u][j];
            }
        }
    }
    if(u == v) return u;
    for(int j = __lg(h[u]);j >= 0; --j){
        if(up[u][j] != up[v][j]){
            u = up[u][j];
            v = up[v][j];
        }
    }
    return up[u][0];
}
int lc[nmax],vis[nmax], c[nmax];
queue<int> q;
vector<int> mau[nmax];
int r[nmax], ans[nmax];

int get(int u){
    return r[u] ? r[u] = get(r[u]) : u;
}
void Union(int u, int v, int val){
    int x = get(u);
    int y = get(v);
    if(x != y){
        r[x] = y;
        lc[y] = lca(lc[y], lc[x]);
        ans[y] += ans[x];
        ans[y] += val;
    }

}
void bfs(int x){
    vis[x] = 1;
    for(auto p : mau[x])q.push(p);
    while(q.size()){
        int u = q.front();q.pop();
        if(u == lc[get(u)]) continue;
        int v = up[u][0];
        if(get(u) == get(v)) continue;
        if(vis[c[v]] == 1) ans[get(u)] = oo;
        else{
            vis[c[v]] = 1;
            for(auto p : mau[c[v]]){
                q.push(p);
                Union(u, p, 1);
            }
        }
    }
}
int idx[nmax];

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);
    }
    dfs(1, -1);
    for(int i = 1; i <= n; ++i){
        cin >> c[i];
        mau[c[i]].push_back(i);
        lc[i] = idx[i] = i;
    }
    for(int i = 1; i <= k; ++i){
        for(int j = 0; j < mau[i].size() - 1; ++j){
            Union(mau[i][j], mau[i][j + 1],0);
        }
    }
    sort(idx + 1, idx + 1 + n, [](int a, int b){
         return h[a] < h[b];
         });

    for(int i = 1; i <= n; ++i){
        int u = idx[i];

        if(vis[c[u]]) continue;
        bfs(c[u]);
    }
    int res = oo;
    for(int i = 1; i <= n; ++i) res = min(res, ans[get(i)]);
    cout << res;


}
/*




*/

Compilation message (stderr)

capital_city.cpp:87:1: warning: ISO C++ forbids declaration of 'main' with no type [-Wreturn-type]
   87 | main(){
      | ^~~~
capital_city.cpp: In function 'int main()':
capital_city.cpp:106:26: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  106 |         for(int j = 0; j < mau[i].size() - 1; ++j){
      |                        ~~^~~~~~~~~~~~~~~~~~~

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