Submission #778764

#	Time	Username	Problem	Language	Result	Execution time	Memory
778764		raysh07	Capital City (JOI20_capital_city)	C++17	100 / 100	290 ms	65032 KiB

capital_city

#include <bits/stdc++.h>
using namespace std;
 
#define sz(v) int(v.size())
#define ar array
typedef long long ll;
const int N = 2e5+10, MOD = 998244353;
 
struct DSU {
    vector<int> par, sz;
    DSU(int n): par(n) {
        iota(par.begin(), par.end(), 0);
        sz.assign(n, 1);
    }
 
    int find_set(int v) {
        return v == par[v] ? v : par[v] = find_set(par[v]);
    }
 
    bool union_sets(int a, int b) {
        if ((a = find_set(a)) == (b = find_set(b))) return false;
        if (sz[a] < sz[b]) swap(a, b);
        par[b] = a, sz[a] += sz[b], sz[b] = 0;
        return true;
    }
};
 
int n, k, col[N], my_out[N], tree_par[N], tree_st[N], tree_en[N], et[N];
vector<int> adj[N];
 
int tt = 0;
void dfs(int c, int p) {
    tree_par[c] = p;
    et[tree_st[c] = tt++] = c;
    for (int nxt : adj[c]) if (nxt != p) {
        dfs(nxt, c);
    }
    tree_en[c] = tt-1;
}
 
int par[N], min_st[N], max_st[N], cnt[N];
set<int> st[N];
multiset<int> ms[N];
 
void merge(int a, int b) {
    a = par[a], b = par[b];
    if (a == b) return;
    if (sz(st[a]) < sz(st[b])) swap(a, b);
 
    for (int x : st[b]) {
        ms[a].erase(x);
        st[a].insert(x);
        par[x] = a;
    }
    st[b].clear();
    cnt[a] += cnt[b];
 
    for (int x : ms[b]) {
        if (!st[a].count(x)) {
            ms[a].insert(x);
        }
    }
    ms[b].clear();
 
    min_st[a] = min(min_st[a], min_st[b]);
    max_st[a] = max(max_st[a], max_st[b]);
}
 
int get_out(int c) {
    c = par[c];
 
    int maybe = et[min_st[c]];
    int undo = -1;
    if (tree_par[maybe] != -1 && tree_en[maybe] >= max_st[c]) {
        undo = col[tree_par[maybe]];
        auto it = ms[c].find(undo);
        if (it != ms[c].end()) {
            ms[c].erase(it);
        } else {
            undo = -1;
        }
    }
 
    int ans = -1;
    if (sz(ms[c])) ans = *ms[c].begin();
 
    if (undo != -1) {
        ms[c].insert(undo);
    }
    return ans;
}
 
 
void solve() {
    cin >> n >> k;
    for (int i = 0; i < n-1; i++) {
        int a, b; cin >> a >> b, --a, --b;
        adj[a].push_back(b), adj[b].push_back(a);
    }
    dfs(0, -1);
 
    for (int i = 0; i < n; i++) cin >> col[i], --col[i];
 
    for (int i = 0; i < k; i++) {
        par[i] = i;
        cnt[i] = 1;
        st[i].insert(i);
    }
 
    fill(min_st, min_st + n, MOD);
    fill(max_st, max_st + n, -1);
    for (int i = 0; i < n; i++) {
        min_st[col[i]] = min(min_st[col[i]], tree_st[i]);
        max_st[col[i]] = max(max_st[col[i]], tree_st[i]);
        if (tree_par[i] != -1 && col[tree_par[i]] != col[i]) {
            ms[col[i]].insert(col[tree_par[i]]);
        }
    }
 
    memset(my_out, -1, sizeof(my_out));
    DSU d(k);
    for (int root = 0; root < k; root++) {
        my_out[par[root]] = get_out(root);
        while (my_out[par[root]] != -1) {
            if (d.union_sets(root, my_out[par[root]])) {
                break;
            } else {
                vector<int> use{par[my_out[par[root]]]};
                while (use.back() != par[root]) {
                    use.push_back(par[my_out[use.back()]]);
                }
 
                use.pop_back();
                for (int x : use) {
                    merge(x, root);
                }
                my_out[par[root]] = get_out(root);
            }
        }
    }
 
    int best = MOD;
    for (int i = 0; i < k; i++) {
        int c = par[i];
        if (my_out[c] == -1) {
            best = min(best, cnt[c] - 1);
        }
    }
 
    cout << best << '\n';
}
 
int main() {
    ios::sync_with_stdio(false); cin.tie(0);
    int T = 1;
    // cin >> T;
    while (T--) solve();
}

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