Submission #1143810

#	Time	Username	Problem	Language	Result	Execution time	Memory
1143810		Persia	Capital City (JOI20_capital_city)	C++17	100 / 100	363 ms	123652 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];

// additional
vector<int> G2[4 * N];
int st[4 * N];
int c1;
vector<int> b[N];

// hld decomposition
int sz[N], big[N], h[N], parr[N];
int in[N], top[N], cnt2;
int a[N];

// scc
int num[4 * N], low[4 * N], cnt;
int col[4 * N], ct;
int val[4 * N];
int bad[4 * N];
bool mark[4 * N];
vector<int> G3[4 * N];
vector<int> stk;

// 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) {
    // cerr << u << " ";

    in[u] = ++cnt2;
    a[cnt2] = 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 addedge(int x, int y) {
    G2[x].push_back(y);

    // cerr << x << " " << y << "\n";
}

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

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

void addedge3(int u, int p, int x) {
    // assert(h[u] >= h[p]);
    while(top[u] != top[p]) {
        // cerr << u << " ";
        addedge2(1, 1, n, in[top[u]], in[u], x);
        u = parr[top[u]];
    }
    // cerr << u << " ";
    if(in[u] > in[p]) swap(u, p);
    addedge2(1, 1, n, in[u], in[p], x);
}

void dfs(int u) {
    num[u] = low[u] = ++cnt;
    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 dp(int u) {
    if(mark[u]) return;
    mark[u] = 1;
    for(int v : G3[u]) {
        dp(v);
        bad[u] |= bad[v];
    }
}

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

    // Dnc Graph + HLD AC

    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();
    c1 = 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][0];
        for(int j : b[i]) u = lca(u, j);
        for(int j : b[i]) {
            addedge3(j, u, i);
            // break;
        }

        // for(int j : b[i]) cout << j << " ";
        // cerr << i << " " << u << "\n";
    }    

    // addedge3(3, 1, 1);

    for(int i = 1; i <= c1; i++) if(!num[i]) dfs(i);
    for(int i = 1; i <= c1; i++) {
        for(int j : G2[i]) if(col[i] != col[j]) {
            bad[col[i]] |= (val[col[j]] > 0);
            G3[col[i]].push_back(col[j]);
        }
    }
    res = n - 1;
    for(int i = 1; i <= ct; i++) if(val[i] > 0) {
        dp(i);
        if(!bad[i]) res = min(res, val[i] - 1);
    }
    cout << res;

    // for(int i = 1; i <= k; i++) cout << col[i] << " ";

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

    // for(int i = 1; i <= k; i++) cerr << col[i] << " ";
    // for(int i = 1; i <= n; i++) cerr << a[i] << " ";
    // for(int i = 1; i <= n; i++) cerr << i << " " << in[i] << "\n";

    return 0 ^ 0;
}

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