Submission #1211182

#	Time	Username	Problem	Language	Result	Execution time	Memory
1211182		dosts	Capital City (JOI20_capital_city)	C++20	100 / 100	322 ms	31012 KiB

capital_city

#include <bits/stdc++.h>
#pragma GCC optimize("O3,unroll-loops")
#pragma GCC target("avx2")
//#define int long long
#define pii pair<int,int>
#define vi vector<int>
#define ff first
#define ss second
#define sp << " " <<
#define all(x) x.begin(),x.end()
#define big(x) ((int)(x.size()))
using namespace std;
const int MOD = 1e9+7, LIM = 1e9, inf = 2e9,MAXN = 1e6+1;

const int N = 2e5+1;
vi edges[N];
int sz[N],dead[N],mark[N],col[N],par[N];
int ans = inf;
int marker = 0;

int szs(int node,int p) {
    sz[node] = 1;
    for (auto it : edges[node]) if (it != p && !dead[it]) sz[node]+=szs(it,node);
    return sz[node];
}

int centro(int node,int p,int s) {
    for (auto it : edges[node]) {
        if (it == p || dead[it]) continue;
        if (sz[it]*2 > s) return centro(it,node,s);
    }
    return node;
}
void dfs(int node,int p) {
    mark[node] = marker;
    par[node] = p;
    for (auto it : edges[node]) {
        if (dead[it] || it == p) continue;
        dfs(it,node);
    }   
}

vi whohas[N];

vi pshd;
vi done(N,0);
vi touched(N,0);
vi toca;
void decompose(int node) {
    ++marker;
    int s = szs(node,node);
    int r = centro(node,node,s);
    dfs(r,r);
    queue<int> q;
    bool cool = 1;
    for (auto it : whohas[col[r]]) {
        if (mark[it] != marker) {
            cool = 0;
            break;
        }
        q.push(it);
    }
    done[col[r]] = 1;
    pshd.push_back(col[r]);
    int hadtodo = 1;
    while (!q.empty() && cool && hadtodo <= ans) {
        int f = q.front();
        q.pop();
        while (f != r && !touched[f]) {
            touched[f] = 1;
            toca.push_back(f);
            if (!done[col[f]]) {
                done[col[f]] = 1;
                pshd.push_back(col[f]);
                hadtodo++;
                if (hadtodo > ans) break;
                for (auto it : whohas[col[f]]) {
                    if (mark[it] != marker) {
                        cool = 0;
                        break;
                    }
                    q.push(it);
                }
            }
            f = par[f];
        }
    }
    if (cool) ans = min(ans,hadtodo-1);
    dead[r] = 1;
    for (auto it : pshd) done[it] = 0;
    pshd.clear();
    for (auto it : toca) touched[it] = 0;
    toca.clear();
    for (auto it : edges[r]) if (!dead[it]) decompose(it);
}

void solve() {
    int n,k;
    cin >> n >> k;
    for (int i=1;i<n;i++) {
        int a,b;
        cin >> a >> b;
        edges[a].push_back(b);
        edges[b].push_back(a);
    }
    for (int i=1;i<=n;i++) cin >> col[i];
    for (int i=1;i<=n;i++) whohas[col[i]].push_back(i);
    decompose(1);
    cout << ans << '\n';
} 

signed main() {
    ios_base::sync_with_stdio(0); cin.tie(0);
    #ifdef Dodi
    freopen("in.txt", "r", stdin);
    freopen("out.txt", "w", stdout);
    #endif
    int t = 1;
    //cin >> t;
    while (t --> 0) solve();
}

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