Submission #1004836

#	Time	Username	Problem	Language	Result	Execution time	Memory
1004836		Unforgettablepl	Capital City (JOI20_capital_city)	C++17	11 / 100	3053 ms	45172 KiB

capital_city

// #pragma GCC optimize("O3")
#include <bits/stdc++.h>
using namespace std;

#define int long long

vector<int> adj[200001];
bool processed[200001];
int siz[200001];
bool current[200001];
int par[200001];
int visitedcol[200001];
int col[200001];
vector<int> nodes[200001];

void set_all(int x,int p,bool set_val){
    if(processed[x])return;
    current[x] = set_val;
    par[x] = p;
    for(int&i:adj[x])if(i!=p)set_all(i,x,set_val);
}

int centroid;

int centroidcalc(int x,int p,int tot){
    if(processed[x])return 0;
    siz[x] = 1;
    bool works = true;
    for(int&i:adj[x])if(i!=p)if(centroidcalc(i,x,tot)>tot/2)works=false;
    if(tot-siz[x]>tot/2)works=false;
    if(works){
        centroid = x;
        return 1e15;
    }
    return siz[x];
}

int calc(int x){
    centroidcalc(x,0,siz[x]);
    set_all(centroid,0,true); // Get them into processing stage
    queue<int> q; // Queue of colours
    q.emplace(col[centroid]);
    int ans = -1;
    bool works = true;
    while(!q.empty()){
        int curr = q.front();q.pop();
        if(visitedcol[curr]==centroid)continue;
        visitedcol[curr]=centroid;
        ans++;
        for(int&i:nodes[curr]){
            if(!current[i]){
                while(!q.empty())q.pop();
                ans = 1e15;
                break;
            }
            if(par[i])q.emplace(col[par[i]]);
        }
    }
    set_all(centroid,0,false);
    processed[centroid]=true;
    for(int&i:adj[centroid])if(!processed[i])ans=min(ans,calc(i));
    return ans;
}

int32_t main(){
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);
    int n,k;
    cin >> n >> k;
    for(int i=1;i<n;i++){
        int a,b;cin>>a>>b;
        adj[a].emplace_back(b);
        adj[b].emplace_back(a);
    }
    for(int i=1;i<=n;i++){
        cin>>col[i];
        nodes[col[i]].emplace_back(i);
    }
    cout << calc(1) << '\n';
}

Compilation message (stderr)

capital_city.cpp: In function 'long long int calc(long long int)':
capital_city.cpp:44:10: warning: unused variable 'works' [-Wunused-variable]
   44 |     bool works = true;
      |          ^~~~~

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