Submission #937521

#	Time	Username	Problem	Language	Result	Execution time	Memory
937521		velislavgarkov	Capital City (JOI20_capital_city)	C++14	100 / 100	446 ms	51560 KiB

capital_city

#include <iostream>
#include <queue>
#include <vector>
using namespace std;
const int MAXN=2e5+10;
const int INF=1e9;
vector <int> diff;
vector <int> v[MAXN], vc[MAXN], ch[MAXN];
int col[MAXN], ans;
int par[MAXN], in[MAXN], out[MAXN], cnt;
int d[MAXN];
bool seen[MAXN], used[MAXN];
int n, k;
void find_ss(int x, int p) {
    d[x]=1;
    for (auto i:v[x]) {
        if (i==p || seen[i]) continue;
        find_ss(i,x);
        d[x]+=d[i];
    }
}
int find_centroid(int x, int ss) {
    for (auto i:v[x]) {
        if (seen[x]) continue;
        if (d[i]>d[x]) continue;
        if (d[i]*2>ss) return find_centroid(i,ss);
    }
    return x;
}
int decompose(int x) {
    find_ss(x,-1);
    int cur=find_centroid(x,d[x]);
    seen[cur]=true;
    for (auto i:v[cur]) {
        if (seen[i]) continue;
        int to=decompose(i);
        ch[cur].push_back(to);
    }
    return cur;
}
void dfs_cent(int x) {
    in[x]=cnt;
    cnt++;
    for (auto i:ch[x]) dfs_cent(i);
    out[x]=cnt;
    cnt++;
}
void dfs(int x, int p) {
    for (auto i:v[x]) {
        if (i==p || seen[i]) continue;
        par[i]=x;
        dfs(i,x);
    }
}
bool in_subtree(int x, int i) {
    return (in[x]<=in[i] && out[x]>=out[i]);
}
void bfs_ans(int start) {
    if (!diff.empty()) diff.clear();
    queue <int> q;
    diff.push_back(col[start]);
    q.push(col[start]);
    used[col[start]]=true;
    int s=0;
    bool fl=false;
    while (!q.empty()) {
        int x=q.front();
        q.pop();
        s++;
        for (auto i:vc[x]) {
            if (!in_subtree(start,i)) {
                fl=true;
                break;
            }
            if (par[i]!=0 && !used[col[par[i]]]) {
                used[col[par[i]]]=true;
                q.push(col[par[i]]);
                diff.push_back(col[par[i]]);
            }
        }
        if (fl) break;
    }
    if (!fl) ans=min(ans,s-1);
    for (auto i:diff) used[i]=false;
}
void find_ans(int x) {
    par[x]=0;
    dfs(x,-1);
    bfs_ans(x);
    seen[x]=true;
    for (auto i:ch[x]) find_ans(i);
}
int main () {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    int a, b;
    cin >> n >> k;
    for (int i=0;i<n-1;i++) {
        cin >> a >> b;
        v[a].push_back(b);
        v[b].push_back(a);
    }
    for (int i=1;i<=n;i++) {
        cin >> col[i];
        vc[col[i]].push_back(i);
    }
    ans=INF;
    int start=decompose(1);
    dfs_cent(start);
    for (int i=1;i<=n;i++) seen[i]=false;
    find_ans(start);
    cout << ans << endl;
    return 0;
}

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