Submission #1138302

#	Time	Username	Problem	Language	Result	Execution time	Memory
1138302		OI_Account	Capital City (JOI20_capital_city)	C++20	11 / 100	2359 ms	552260 KiB

capital_city

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

const int N = 200'000;
const int S = N * 19;
const int M = 17;

int n, k, c[N + 10], ver;
int h[N + 10], dp[N + 10][M + 2], idSp[N + 10][M + 2];
int cntCmp, id[S + 10], mn[S + 10], mx[S + 10];
vector<int> adj[N + 10], vec, pack[N + 10];
vector<int> out[S + 10], in[S + 10];
int sum[S + 10];
bool notOut[S + 10];

void readInput() {
    cin >> n >> k; 
    for (int i = 1; i < n; i++) {
        int u, v;
        cin >> u >> v;
        adj[u].push_back(v);
        adj[v].push_back(u);
    }
    for (int i = 1; i <= n; i++) {
        cin >> c[i];
        pack[c[i]].push_back(i);
    }
}

void calcIdSp() {
    int pnt = k;
    for (int i = 1; i <= n; i++)
        for (int j = 0; j <= M; j++) {
            idSp[i][j] = ++pnt;
            out[pnt].reserve(2);
        }
    ver = pnt;
}

void addEdge(int u, int v) {
    out[u].push_back(v);
    in[v].push_back(u);
}

void dfs(int u = 1, int par = 0) {
    h[u] = h[par] + 1;
    dp[u][0] = par;
    addEdge(idSp[u][0], c[u]);
    for (int j = 1; j <= M && dp[u][j - 1]; j++) {
        dp[u][j] = dp[dp[u][j - 1]][j - 1];
        addEdge(idSp[u][j], idSp[u][j - 1]);
        addEdge(idSp[u][j], idSp[dp[u][j - 1]][j - 1]);
    }
    for (auto v: adj[u])
        if (v != par)
            dfs(v, u);
}

int LCA(int u, int v) {
    if (h[u] < h[v])
        swap(u, v);
    for (int j = M; j >= 0; j--)
        if (h[u] - h[v] >= (1 << j))
            u = dp[u][j];
    if (u == v)
        return u;
    for (int j = M; j >= 0; j--)
        if (dp[u][j] != dp[v][j]) {
            u = dp[u][j];
            v = dp[v][j];
        }
    return dp[u][0];
}

void addEdges(int x, int u, int v) {
    /*while (h[u] - h[v] >= 2) {
        addEdge(x, idSp[u][1]);
        u = dp[u][1];
    }*/
    /*while (u != v) {
        addEdge(x, idSp[u][0]);
        u = dp[u][0];
    }*/
    for (int j = M; j >= 0; j--)
        if (h[u] - h[v] >= (1 << j)) {
            addEdge(x, idSp[u][j]);
            u = dp[u][j];
        }
    addEdge(x, c[v]);
}

void calcGraph() {
    for (int i = 1; i <= k; i++) {
        int lca = pack[i][0];
        for (int j = 1; j < pack[i].size(); j++)
            lca = LCA(lca, pack[i][j]);
        for (auto u: pack[i])
            addEdges(c[u], u, lca);
    }
}

inline void dfsOut(int u) {
    id[u] = -1;
    for (auto v: out[u])
        if (!id[v])
            dfsOut(v);
    vec.push_back(u);
}

inline void dfsIn(int u) {
    id[u] = cntCmp;
    for (auto v: in[u])
        if (id[v] == -1)
            dfsIn(v);
}

void SCC() {    
    vec.reserve(ver);
    for (int i = 1; i <= ver; i++)
        if (!id[i])
            dfsOut(i);
    reverse(vec.begin(), vec.end());
    for (auto u: vec)
        if (id[u] == -1) {
            cntCmp++;
            dfsIn(u);
        }
}

void dfsMinMax(int u = 1, int par = 0) {
    mn[idSp[u][0]] = mx[idSp[u][0]] = id[c[u]];
    for (int j = 1; j <= M && dp[u][j - 1]; j++) {
        mn[idSp[u][j]] = min(mn[idSp[u][j - 1]], mn[idSp[dp[u][j - 1]][j - 1]]);
        mx[idSp[u][j]] = max(mx[idSp[u][j - 1]], mx[idSp[dp[u][j - 1]][j - 1]]);
    }
    for (auto v: adj[u])
        if (v != par)
            dfsMinMax(v, u);
}

void init() {
    for (int u = 1; u <= k; u++) {
        sum[id[u]]++;
        for (auto v: out[u])
            if (v <= k)
                notOut[id[u]] |= (id[v] != id[u]);
            else
                notOut[id[u]] |= (mn[v] != id[u] || mx[v] != id[u]);
    }
}

void calcAns() {
    int ans = k;
    for (int i = 1; i <= cntCmp; i++)
        if (notOut[i] == false && sum[i])
            ans = min(ans, sum[i] - 1);
    cout << ans << flush;
}

int main() {
    ios_base::sync_with_stdio(false); cin.tie(0); cout.tie(0);
    readInput();
    calcIdSp();
    dfs();
    calcGraph();
    SCC();
    dfsMinMax();
    init();
    calcAns();
    return 0;
}

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