Submission #537358

#	Time	Username	Problem	Language	Result	Execution time	Memory
537358		wiwiho	Capital City (JOI20_capital_city)	C++14	31 / 100	690 ms	76180 KiB

capital_city

#include <bits/stdc++.h>

#define mp make_pair
#define F first
#define S second
#define eb emplace_back
#define iter(a) a.begin(), a.end()
#define lsort(a) sort(iter(a))
#define gsort(a) sort(iter(a), greater<>())
#define printv(a, b) { \
    for(auto pv : a) b << pv << " "; \
    b << "\n"; \
}

using namespace std;

using pii = pair<int, int>;

int n, m;
vector<int> C;
vector<vector<int>> g;
vector<vector<int>> cv;

vector<int> pr, in, out, dpt;
vector<vector<int>> anc;

vector<int> dsu, mx, sz, top;
vector<set<int>> dc;

int ts = 0;
const int SZ = 20;

void init(){
    C.resize(n + 1);
    g.resize(n + 1);
    cv.resize(m + 1);

    pr.resize(n + 1);
    in.resize(n + 1);
    out.resize(n + 1);
    anc.resize(SZ, vector<int>(n + 1));
    dpt.resize(n + 1);

    dsu.resize(n + 1);
    iota(iter(dsu), 0);
    mx.resize(n + 1);
    dc.resize(n + 1);
    sz.resize(n + 1, 1);
    top.resize(n + 1);
    iota(iter(top), 0);
}

int findDSU(int a){
    if(dsu[a] != a) dsu[a] = findDSU(dsu[a]);
    return dsu[a];
}

void unionDSU(int a, int b){
    a = findDSU(a);
    b = findDSU(b);
    if(a == b) return;
    if(sz[a] < sz[b]) swap(a, b);
    dsu[b] = a;
    sz[a] += sz[b];
    for(int i : dc[b]) dc[a].insert(i);
    if(in[top[b]] < in[top[a]]) top[a] = top[b];
    mx[a] = max(mx[a], mx[b]);
}

void dfs(int now, int p){
    pr[now] = p;
    in[now] = ts++;
    dpt[now] = dpt[p] + 1;
    anc[0][now] = p;
    for(int i : g[now]){
        if(i == p) continue;
        dfs(i, now);
    }
    out[now] = ts++;
}

bool isAnc(int a, int b){
    return in[a] <= in[b] && out[a] >= out[b];
}

int lca(int a, int b){
    if(isAnc(a, b)) return a;
    for(int i = SZ - 1; i >= 0; i--){
        if(!isAnc(anc[i][a], b)) a = anc[i][a];
    }
    return anc[0][a];
}

int dis(int a, int b){
    int l = lca(a, b);
    return dpt[a] + dpt[b] - 2 * dpt[l];
}

int getsz(int c){
    sort(iter(cv[c]), [&](int x, int y){ return in[x] < in[y]; });
    int ans = 0;
    int rt = cv[c][0];
    for(int i : cv[c]) rt = lca(rt, i);
    int lst = rt;
    for(int i : cv[c]){
        ans += dis(lst, i);
        lst = i;
    }
    ans += dis(lst, rt);
    ans /= 2; ans++;
    return ans;
}

int main(){
    ios_base::sync_with_stdio(false);
    cin.tie(0);

    cin >> n >> m;
    init();

    for(int i = 1; i < n; i++){
        int u, v;
        cin >> u >> v;
        g[u].eb(v);
        g[v].eb(u);
    }
    for(int i = 1; i <= n; i++){
        cin >> C[i];
        cv[C[i]].eb(i);
        dc[i].insert(C[i]);
    }

    dfs(1, 1);
    for(int i = 1; i < SZ; i++){
        for(int j = 1; j <= n; j++){
            anc[i][j] = anc[i - 1][anc[i - 1][j]];
        }
    }

    vector<int> csz(m + 1);
    for(int i = 1; i <= m; i++) csz[i] = getsz(i);
    
    vector<int> owo(m);
    iota(iter(owo), 1);
    sort(iter(owo), [&](int x, int y){ return csz[x] < csz[y]; });
    vector<int> id(m + 1);
    for(int i = 0; i < m; i++) id[owo[i]] = i;
    for(int i = 1; i <= n; i++){
        mx[i] = id[C[i]];
    }

    //printv(owo, cerr);

    int ans = n;
    for(int i = 0; i < m; i++){
        int c = owo[i];
        //cerr << "do " << c << "\n";
        int rt = cv[c][0];
        for(int i : cv[c]) rt = lca(rt, i);
        
        for(int v : cv[c]){
            while(!isAnc(v, rt)){
                unionDSU(v, pr[v]);
                v = top[findDSU(v)];
            }
        }
        //for(int i = 1; i <= n; i++) cerr << findDSU(i) << " ";
        //cerr << "\n";
        rt = findDSU(rt);
        //cerr << "ok " << rt << " " << dc[rt].size() << " " << mx[rt] << "\n";
        if(mx[rt] == i) ans = min(ans, (int)dc[rt].size() - 1);
    }
    cout << ans << "\n";

    return 0;
}

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