Submission #898487

#	Time	Username	Problem	Language	Result	Execution time	Memory
898487		box	Capital City (JOI20_capital_city)	C++17	100 / 100	762 ms	149900 KiB

capital_city

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

#ifdef LOCAL
#define bug(x) cerr << "L" << __LINE__ << " | " << #x << " = " << x << endl;
#else
#define bug(x)
#define cerr if (0) cerr
#endif

#define ar array
#define sz(v) int(std::size(v))
#define all(v) std::begin(v), std::end(v)
using i64 = long long;
using vi = vector<int>;
using pi = pair<int, int>;

const int MAXN = 2e5;

int N, K, C[MAXN];
vi adj[MAXN], who[MAXN];
int num;
int cnt[MAXN], child[MAXN], parent[MAXN], depth[MAXN], root[MAXN];

const int MAXS = MAXN * 3;

vi g[MAXS], rg[MAXS], order;
int id[MAXS], colors[MAXN];

void dfs_ord(int i) {
    id[i] = 0;
    for (int j : g[i]) if (id[j] == -1) dfs_ord(j);
    order.push_back(i);
}

void dfs_id(int i, int c) {
    id[i] = c;
    for (int j : rg[i]) if (id[j] == -1) dfs_id(j, c);
}

void make_edge(int i, int j) {
    g[i].push_back(j);
    rg[j].push_back(i);
}

struct segt {
    int l, r;
    segt *lc = NULL, *rc = NULL;
    int id;
    segt(int l, int r) : l(l), r(r) {
        id = num++;
        if (l < r) {
            int m = (l + r) / 2;
            lc = new segt(l, m);
            rc = new segt(m + 1, r);
            make_edge(lc->id, id);
            make_edge(rc->id, id);
        }
    }
    void qry(int ql, int qr, vi& v) {
        if (ql <= l && qr >= r) v.push_back(id);
        else {
            if (ql <= lc->r) lc->qry(ql, qr, v);
            if (qr >= rc->l) rc->qry(ql, qr, v);
        }
    }
} *ids[MAXN];

vi qry(int i, int l, int r) {
    vi v;
    ids[i]->qry(l, r, v);
    return v;
}

int dfs_sz(int i) {
    int s = 1, x = 0;
    child[i] = -1;
    for (int j : adj[i]) if (parent[i] != j) {
        parent[j] = i;
        depth[j] = depth[i] + 1;
        int y = dfs_sz(j);
        if (y > x) x = y, child[i] = j;
        s += y;
    }
    return s;
}

void dfs_hvy(int r, int i) {
    cnt[root[i] = r]++;
    if (~child[i]) dfs_hvy(r, child[i]);
    for (int j : adj[i]) if (parent[i] != j && child[i] != j) dfs_hvy(j, j);
}

int lca(int i, int j) {
    while (root[i] != root[j]) {
        if (depth[root[i]] < depth[root[j]]) swap(i, j);
        i = parent[root[i]];
    }
    return depth[i] < depth[j] ? i : j;
}

void make_path(int p, int i, vi& v) {
    while (root[i] != root[p]) {
        int r = root[i];
        ids[r]->qry(0, depth[i] - depth[r], v);
        i = parent[r];
    }
    int r = root[i];
    ids[r]->qry(depth[p] - depth[r], depth[i] - depth[r], v);
}

int main() {
    ios::sync_with_stdio(0); cin.tie(0);
    cin >> N >> K;
    for (int h = 0; h < N - 1; h++) {
        int i, j; cin >> i >> j, i--, j--;
        adj[i].push_back(j), adj[j].push_back(i);
    }
    for (int i = 0; i < N; i++) {
        cin >> C[i], C[i]--;
        who[C[i]].push_back(i);
    }
    num = K;
    parent[0] = -1;
    dfs_sz(0);
    dfs_hvy(0, 0);
    for (int i = 0; i < N; i++) if (root[i] == i) ids[i] = new segt(0, cnt[i] - 1);
    assert(num <= MAXS);
    for (int i = 0; i < N; i++) {
        int r = root[i];
        for (int v : qry(r, depth[i] - depth[r], depth[i] - depth[r]))
            make_edge(C[i], v);
    }
    for (int c = 0; c < K; c++) {
        int top = who[c].at(0);
        for (int i : who[c]) top = lca(top, i);
        vi v;
        for (int i : who[c]) make_path(top, i, v);
        for (int u : v) make_edge(u, c);
    }
    fill(id, id + num, -1);
    for (int i = 0; i < num; i++) if (id[i] == -1) dfs_ord(i);
    reverse(all(order));
    fill(id, id + num, -1);
    int c = 0;
    for (int i : order) if (id[i] == -1) dfs_id(i, c++);
    for (int c = 0; c < K; c++) colors[id[c]]++;
    for (int i = 0; i < num; i++) for (int j : g[i]) if (id[i] != id[j]) colors[id[j]] = INT_MAX;
    int ans = INT_MAX;
    for (int c = 0; c < K; c++) ans = min(ans, colors[id[c]] - 1);
    cout << ans << '\n';

}

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