Submission #1244891

#	Time	Username	Problem	Language	Result	Execution time	Memory
1244891		quannnguyen2009	Capital City (JOI20_capital_city)	C++20	11 / 100	2145 ms	589824 KiB

capital_city

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

// naive SCC + Tarjan solution without HLD/segment tree optimization
// vertices are colors 1..k
// graph edges: u->v if color u depends on color v along tree paths

static const int MAXN = 200000;
static const int LOG = 18;

int n, k;
vector<int> tree[MAXN+1];
int color[MAXN+1];
int parent[LOG][MAXN+1], depth[MAXN+1];

// DFS to fill parent[0] and depth
void dfs(int u, int p) {
    parent[0][u] = p;
    for (int v : tree[u]) {
        if (v == p) continue;
        depth[v] = depth[u] + 1;
        dfs(v, u);
    }
}

// Binary-lifted LCA
int lca(int u, int v) {
    if (depth[u] < depth[v]) swap(u, v);
    int diff = depth[u] - depth[v];
    for (int i = 0; i < LOG; i++)
        if (diff & (1 << i)) u = parent[i][u];
    if (u == v) return u;
    for (int i = LOG-1; i >= 0; i--) {
        if (parent[i][u] != parent[i][v]) {
            u = parent[i][u];
            v = parent[i][v];
        }
    }
    return parent[0][u];
}

// dependency graph on colors
vector<vector<int>> G2;

// Tarjan's SCC data
vector<int> dfn, low, comp;
vector<bool> onstack;
stack<int> st;
int timer = 0, compCnt = 0;

void tarjan(int u) {
    dfn[u] = low[u] = ++timer;
    st.push(u);
    onstack[u] = true;
    for (int v : G2[u]) {
        if (!dfn[v]) {
            tarjan(v);
            low[u] = min(low[u], low[v]);
        } else if (onstack[v]) {
            low[u] = min(low[u], dfn[v]);
        }
    }
    if (low[u] == dfn[u]) {
        compCnt++;
        while (true) {
            int v = st.top(); st.pop();
            onstack[v] = false;
            comp[v] = compCnt;
            if (v == u) break;
        }
    }
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(NULL);

    cin >> n >> k;
    for (int i = 1; i < n; i++) {
        int u, v; cin >> u >> v;
        tree[u].push_back(v);
        tree[v].push_back(u);
    }
    for (int i = 1; i <= n; i++) {
        cin >> color[i];
    }

    // build depth & parent table
    depth[1] = 0;
    dfs(1, 0);
    for (int i = 1; i < LOG; i++) {
        for (int u = 1; u <= n; u++) {
            parent[i][u] = parent[i-1][ parent[i-1][u] ];
        }
    }

    // group nodes by color
    vector<vector<int>> byColor(k+1);
    for (int u = 1; u <= n; u++) byColor[color[u]].push_back(u);

    // init dependency graph for colors 1..k
    G2.assign(k+1, {});

    // build naive edges
    for (int c = 1; c <= k; c++) {
        if (byColor[c].empty()) continue;
        // compute LCA of all nodes of this color
        int u = byColor[c][0];
        for (int v : byColor[c]) u = lca(u, v);
        // walk each node up to u
        for (int v : byColor[c]) {
            int x = v;
            while (x != u) {
                int p = parent[0][x];
                G2[c].push_back(color[p]);
                x = p;
            }
        }
    }

    // prepare Tarjan
    dfn.assign(k+1, 0);
    low.assign(k+1, 0);
    comp.assign(k+1, 0);
    onstack.assign(k+1, false);

    // run Tarjan on each color node
    for (int i = 1; i <= k; i++) if (!dfn[i]) tarjan(i);

    // count sizes of components
    vector<int> compSize(compCnt+1, 0);
    for (int i = 1; i <= k; i++) compSize[ comp[i] ]++;

    // build condensed DAG and mark bad components
    vector<vector<int>> dag(compCnt+1);
    vector<bool> bad(compCnt+1, false);
    for (int u = 1; u <= k; u++) {
        for (int v : G2[u]) {
            if (comp[u] != comp[v]) {
                dag[ comp[u] ].push_back(comp[v]);
                if (compSize[ comp[v] ] > 0) bad[ comp[u] ] = true;
            }
        }
    }

    // propagate bad flag
    vector<bool> visited(compCnt+1, false);
    function<void(int)> dfs2 = [&](int u) {
        if (visited[u]) return;
        visited[u] = true;
        for (int v : dag[u]) {
            dfs2(v);
            if (bad[v]) bad[u] = true;
        }
    };
    for (int i = 1; i <= compCnt; i++) dfs2(i);

    // compute answer
    int answer = k - 1;
    for (int i = 1; i <= compCnt; i++) {
        if (!bad[i] && compSize[i] > 0) {
            answer = min(answer, compSize[i] - 1);
        }
    }
    cout << answer;
    return 0;
}

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