Submission #448245

#	Time	Username	Problem	Language	Result	Execution time	Memory
448245		JovanB	Capital City (JOI20_capital_city)	C++17	100 / 100	1545 ms	41932 KiB

capital_city

#include <bits/stdc++.h>
using namespace std;

using ll = long long;
using ld = long double;

int res;

const int N = 200000;

int par[N+5];
vector <int> graf[N+5];
int sz[N+5];
bool erased[N+5];
bool visited[N+5];
bool viscol[N+5];
int col[N+5];
vector <int> vec[N+5];
int szc[N+5];

void dfs_size(int v, int p){
    sz[v] = 1;
    par[v] = p;
    vec[col[v]].push_back(v);
    for(auto c : graf[v]){
        if(erased[c] || c == p) continue;
        dfs_size(c, v);
        sz[v] += sz[c];
    }
}

void dfs_clear(int v, int p){
    visited[v] = 0;
    viscol[col[v]] = 0;
    vec[col[v]].clear();
    for(auto c : graf[v]) if(!erased[c] && c != p) dfs_clear(c, v);
}

int find_centroid(int v, int p, int svi){
    for(auto c : graf[v]) if(!erased[c] && c != p && 2*sz[c] > svi) return find_centroid(c, v, svi);
    return v;
}

void decompose(){
    queue <int> q;
    q.push(1);
    while(!q.empty()){
        int v = q.front();
        q.pop();
        dfs_size(v, 0);
        v = find_centroid(v, 0, sz[v]);
        dfs_clear(v, 0);
        dfs_size(v, 0);
        queue <int> colors;
        colors.push(col[v]);
        viscol[col[v]] = 1;
        visited[v] = 1;
        int cnt = 0;
        while(!colors.empty()){
            int cl = colors.front();
            colors.pop();
            cnt++;
            if(szc[cl] != vec[cl].size()) cnt = N;
            for(auto c : vec[cl]){
                int x = c;
                while(!visited[x]){
                    visited[x] = 1;
                    if(!viscol[col[x]]){
                        viscol[col[x]] = 1;
                        colors.push(col[x]);
                    }
                    x = par[x];
                }
            }
        }
        res = min(res, cnt - 1);
        erased[v] = 1;
        for(auto c : graf[v]) if(!erased[c]) q.push(c);
    }
}

int main(){
    ios_base::sync_with_stdio(false), cin.tie(0);
    cout.precision(10);
    cout << fixed;

    int n, k;
    cin >> n >> k;
    for(int i=1; i<n; i++){
        int a, b;
        cin >> a >> b;
        graf[a].push_back(b);
        graf[b].push_back(a);
    }
    for(int i=1; i<=n; i++){
        cin >> col[i];
        szc[col[i]]++;
    }
    res = k - 1;
    decompose();
    cout << res << "\n";
    return 0;
}

Compilation message (stderr)

capital_city.cpp: In function 'void decompose()':
capital_city.cpp:63:24: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   63 |             if(szc[cl] != vec[cl].size()) cnt = N;
      |                ~~~~~~~~^~~~~~~~~~~~~~~~~

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