Submission #577515

#	Time	Username	Problem	Language	Result	Execution time	Memory
577515		balbit	Capital City (JOI20_capital_city)	C++14	100 / 100	847 ms	43888 KiB

capital_city

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

#define ll long long
#define pii pair<ll, ll>
#define f first
#define s second

#define SZ(x) (int)((x).size())
#define ALL(x) (x).begin(), (x).end()
#define pb push_back

#define MX(a,b) a =  max(a,b)
#define MN(a,b) a =  min(a,b)
#define FOR(i,a,b) for (int i = a; i<b; ++i)
#define REP(i,n) for (int i = 0; i<n; ++i)
#define REP1(i,n) for (int i = 1; i<=n; ++i)

#ifdef BALBIT
#define bug(...) cerr<<"#"<<__LINE__<<"- "<<#__VA_ARGS__<<": ", _do(__VA_ARGS__)
template<typename T> void _do(T && x){cerr<<x<<endl;}
template<typename T, typename ...S> void _do(T && x, S && ...y){cerr<<x<<", "; _do(y...);}
#else
#define bug(...)
#define endl '\n'
#endif // BALBIT


const int maxn = 2e5+5;
const int mod = 1e9+7;


bool dead[maxn];
bool colorok[maxn];
int color[maxn];
int sz[maxn];

vector<int> g[maxn];
vector<int> where[maxn]; // where are the nodes of this color?

void predfs(int v, int p) {
    sz[v] = 1;
    for (int u : g[v]) {
        if (dead[u]) continue;
        if (u != p) {
            predfs(u,v);
            sz[v] += sz[u];
        }
    }
}
int allsz = -1;
int getcen(int v, int p) {
    for (int u : g[v]) {
        if (dead[u]) continue;
        if (u != p && sz[u] > (allsz)/2) {
            return getcen(u,v);
        }
    }
    return v;
}

int dep[maxn];
bool done[maxn]; // make sure to clear!!!
bool inque[maxn]; // make sure to clear!!!
int par[maxn];
vector<int> colorshere;

int subtree[maxn]; // which subtree is it?????

void dfs(int v, int p) {
    done[v] = 0; inque[color[v]] = 0;
    colorshere.pb(color[v]);
    for (int u : g[v]) {
        if (dead[u]) continue;
        if (u!=p) {
            dep[u] = dep[v] + 1; par[u] = v;
            if (p == -1) {
                subtree[u] = u;
            }else{
                subtree[u] = subtree[v];
            }
            dfs(u,v);
        }
    }
}

int re = maxn;

void decomp(int stv, int lay) {
    predfs(stv, -1); allsz = sz[stv];
    int cen = getcen(stv, -1);

    bug(stv, lay, cen);

    stv = -1;

    dep[cen] = 0;
    colorshere.clear();

    dfs(cen, -1);

    sort(ALL(colorshere)); colorshere.resize(unique(ALL(colorshere)) - colorshere.begin());
    done[cen] = 1;

    if (colorok[color[cen]]) {
        queue<int> q; q.push(color[cen]);
        inque[color[cen]] = 1;
        int nmerge = 0;
        bool bad = 0;
        while (!q.empty()) {
            int cc = q.front(); q.pop();
            for (int x : where[cc]) {
                for (int at = x; !done[at]; at= par[at]) {
                    done[at] = 1;
                    int c = color[at];
                    if (!inque[c] && colorok[c]) {
                        inque[c] = 1; q.push(c); ++nmerge;
                    }
                    if (!colorok[c]) {
                        bad = 1; break;
                    }
                }
                if (bad) break;
            }
        }

        if (!bad) {
            bug(nmerge);
            MN(re, nmerge);
        }
    }

    dead[cen] = 1; colorok[color[cen]] = 0;

    for (int t : colorshere) {
        if (colorok[t]) {
            int mn = maxn, mx = -1;
            for (int x : where[t]) {
                MN(mn, subtree[x]); MX(mx, subtree[x]);
            }
            if (mn < mx) {
                colorok[t] = 0;
            }
        }
    }

    for (int u : g[cen]) {
        if (!dead[u]) {
            decomp(u, lay+1);
        }
    }
}

signed main(){
    ios::sync_with_stdio(0), cin.tie(0);
    bug(1,2);

    int n,c; cin>>n>>c;

    REP(i,n-1) {
        int a,b; cin>>a>>b; --a; --b;
        g[a].pb(b); g[b].pb(a);
    }
    REP(i,n) {
        cin>>color[i]; colorok[color[i]] = 1;
        where[color[i]].pb(i);
    }

    decomp(0, 0);

    cout<<re<<endl;

}

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