Submission #1313697

#	Time	Username	Problem	Language	Result	Execution time	Memory
1313697		vedchoudhary	Capital City (JOI20_capital_city)	C++20	100 / 100	1012 ms	56840 KiB

capital_city

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

using namespace std;

using ll = long long;
#define int ll

vector<vector<int>> adj;
vector<int> p;
vector<bool> rdy;
vector<int> sz;
vector<int> sub;
vector<int> cnt;
vector<int> seen;
vector<int> cityOf;
vector<vector<int>> towns;
int ans;

void dfs(int curr, int prev) {
    p[curr] = prev;
    sub[curr] = 1;
    for(int nxt : adj[curr]) {
        if(!rdy[nxt]||nxt==prev) continue;
        dfs(nxt,curr);
        sub[curr] += sub[nxt];
    }
}

int centroid(int curr, int prev, int n) {
    for(int nxt : adj[curr]) {
        if(!rdy[nxt]||nxt==prev) continue;
        if(sub[nxt]>n/2) return centroid(nxt,curr,n);
    }
    return curr;
}

void solve(int c) {
    dfs(c,c);
    c = centroid(c,c,sub[c]);
    dfs(c,c);
    stack<pair<int,int>> s; s.push({c,c});
    set<int> reset;
    while(!s.empty()) {
        auto [curr,prev] = s.top(); s.pop();
        cnt[cityOf[curr]]++;
        reset.insert(cityOf[curr]);
        for(int nxt : adj[curr]) {
            if(!rdy[nxt]||nxt==prev) continue;
            s.push({nxt,curr});
        }
    }
    int cc = 0;
    stack<int> cs;
    cs.push(cityOf[c]);
    while(!cs.empty()) {
        int curr = cs.top(); cs.pop();
        if(seen[curr]) continue;
        if(cnt[curr]!=sz[curr]) { cc = 1000000000; break; }
        seen[curr] = true;
        cc++;
        for(int t : towns[curr]) {
            cs.push(cityOf[p[t]]);
        }
    }
    for(int cty : reset) {
        cnt[cty] = 0;
        seen[cty] = false;
    }
    if(ans==-1) ans = cc;
    else ans = min(ans,cc);
    rdy[c] = false;
    for(int nxt : adj[c]) {
        if(!rdy[nxt]) continue;
        solve(nxt);
    }
}


signed main() {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);
    
    ans = -1;
    int n,k; cin >> n >> k;
    adj.resize(n+1);
    p.resize(n+1);
    rdy.resize(n+1,true);
    sz.resize(k+1);
    sub.resize(n+1);
    cnt.resize(k+1);
    cityOf.resize(n+1);
    seen.resize(k+1);
    towns.resize(k+1);
    for(int i = 0; i < n-1; i++) {
        int a,b; cin >> a >> b;
        adj[a].push_back(b);
        adj[b].push_back(a);
    }

    for(int i = 1; i <= n; i++) {
        cin >> cityOf[i];
        sz[cityOf[i]]++;
        towns[cityOf[i]].push_back(i);
    }

    solve(1);
    cout << ((ans==1||ans==1000000000)?0:ans-1);
    return 0;
}

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