Submission #1143860

#	Time	Username	Problem	Language	Result	Execution time	Memory
1143860		Persia	Capital City (JOI20_capital_city)	C++17	100 / 100	350 ms	123692 KiB

capital_city

#include <bits/stdc++.h>

using namespace std;

#define bit(i, x) (x >> i & 1)
#define ll long long
#define sz(x) (int)x.size()

const int N = 2e5 + 5;
const int mod = 998244353;

mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());
ll rnd(ll l, ll r) {
    return uniform_int_distribution<ll>(l, r)(rng);
}

int n, k;
vector<int> G[N];
int c[N];

// HLD DECOMPOSITION
int sz[N], big[N], parr[N], h[N];
int in[N], top[N], a[N], cnt;

// DNC GRAPH
int st[4 * N], cnt2;
vector<int> G2[4 * N];

// SCC
int num[4 * N], low[4 * N], cnt3;
int col[4 * N], val[4 * N], ct;
vector<int> stk;

// ADDITIONAL
vector<int> b[N];

// DAG
vector<int> G3[4 * N];
int bad[4 * N];
bool mark[4 * N];

// RESULT
int res;

void predfs(int u = 1, int par = -1) {
    sz[u] = 1, big[u] = 0, parr[u] = par;
    for(int v : G[u]) if(v != par) {
        h[v] = h[u] + 1;
        predfs(v, u);
        sz[u] += sz[v];
        if(sz[big[u]] < sz[v]) big[u] = v;
    }
}

void dfshld(int u = 1, int topp = 1) {
    in[u] = ++cnt;
    a[cnt] = c[u];
    top[u] = topp;
    if(big[u]) dfshld(big[u], topp);
    for(int v : G[u]) if(v != parr[u] && v != big[u]) dfshld(v, v);
}

int lca(int x, int y) {
    while(top[x] != top[y]) {
        if(h[top[x]] < h[top[y]]) swap(x, y);
        x = parr[top[x]];
    }
    if(h[x] < h[y]) swap(x, y);
    return y;
}

void build(int id, int l, int r) {
    st[id] = ++cnt2;
    if(l == r) {
        G2[st[id]].push_back(a[l]);
        return;
    }
    int mid = (l + r) / 2;
    build(id * 2, l, mid);
    build(id * 2 + 1, mid + 1, r);
    G2[st[id]].push_back(st[id * 2]);
    G2[st[id]].push_back(st[id * 2 + 1]);
}

void addedge(int id, int l, int r, int u, int v, int x) {
    if(l > v || r < u) return;
    if(l >= u && r <= v) {
        G2[x].push_back(st[id]);
        return;
    }
    int mid = (l + r) / 2;
    addedge(id * 2, l, mid, u, v, x);
    addedge(id * 2 + 1, mid + 1, r, u, v, x);
}

void addedge2(int x, int p) {
    assert(h[x] >= h[p]);
    int u = c[x];
    while(top[x] != top[p]) {
        addedge(1, 1, n, in[top[x]], in[x], u);
        x = parr[top[x]];
    }
    if(in[p] > in[x]) swap(x, p);
    addedge(1, 1, n, in[p], in[x], u);
}

void dfs(int u) {
    num[u] = low[u] = ++cnt3;
    stk.push_back(u);
    for(int v : G2[u]) if(!col[v]) {
        if(num[v]) low[u] = min(low[u], num[v]);
        else {
            dfs(v);
            low[u] = min(low[u], low[v]);
        }
    }
    if(num[u] == low[u]) {
        ct++;
        while(1) {
            int top = stk.back();
            stk.pop_back();
            col[top] = ct;
            val[ct] += (top <= k);
            if(top == u) break;
        }
    }
}

void dag(int u) {
    if(mark[u]) return;
    mark[u] = 1;
    for(int v : G3[u]) {
        dag(v);
        bad[u] |= bad[v];
    }
}

signed main(int argc, char* argv[]) {
    ios_base::sync_with_stdio(0);
    cin.tie(0); cout.tie(0);

    cin >> n >> k;
    for(int i = 1; i < n; i++) {
        int x, y; cin >> x >> y;
        G[x].push_back(y);
        G[y].push_back(x);
    }
    for(int i = 1; i <= n; i++) cin >> c[i];

    predfs();
    dfshld();
    cnt2 = k;
    build(1, 1, n);
    for(int i = 1; i <= n; i++) b[c[i]].push_back(i);
    for(int i = 1; i <= k; i++) if(!b[i].empty()) {
        int u = b[i].front();
        for(int j : b[i]) u = lca(u, j);
        for(int j : b[i]) addedge2(j, u);
    }
    for(int i = 1; i <= cnt2; i++) if(!num[i]) dfs(i);
    for(int i = 1; i <= cnt2; i++) {
        for(int j : G2[i]) if(col[i] != col[j]) {
            G3[col[i]].push_back(col[j]);
            bad[col[i]] |= (val[col[j]] > 0);
        }
    }
    res = k - 1;
    for(int i = 1; i <= ct; i++) {
        dag(i);
        if(!bad[i] && val[i]) res = min(res, val[i] - 1);
    }
    cout << res;

    // for(int i = 1; i <= n; i++) cout << i << " " << in[i] << "\n";
    // for(int i = 1; i <= n; i++) cout << i << " " << big[i] << "\n";
    // for(int i = 1; i <= n; i++) {
    //     for(int j = i; j <= n; j++) {
    //         cout << i << " " << j << " " << lca(i, j) << "\n";
    //     }
    // }

    return 0 ^ 0;
}

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