Submission #808587

#	Time	Username	Problem	Language	Result	Execution time	Memory
808587		FedShat	Capital City (JOI20_capital_city)	C++17	100 / 100	1872 ms	346788 KiB

capital_city

#include <bits/stdc++.h>

using namespace std;
using ll = long long;
constexpr int INF = numeric_limits<int>::max() / 2;
constexpr ll INFLL = numeric_limits<ll>::max() / 2;

template<class T>
istream &operator>>(istream &is, vector<T> &a) {
    for (auto &i : a) {
        is >> i;
    }
    return is;
}

#ifdef __APPLE__
#include "debug.h"
#else
#define debug(...) 42
#endif

struct LCA {
    vector<int> depth, parent, up;

    LCA() {}
    LCA(int n) {
        depth.resize(n);
        parent.resize(n);
        up.resize(n);
    }

    void add_leaf(int v, int par) {
        parent[v] = par;
        depth[v] = depth[par] + 1;
        if (depth[par] - depth[up[par]] == depth[up[par]] - depth[up[up[par]]]) {
            up[v] = up[up[par]];
        } else {
            up[v] = par;
        }
    }

    int la(int v, int k) {
        int h = max(depth[v] - k, 0);
        while (depth[v] != h) {
            if (depth[up[v]] >= h) {
                v = up[v];
            } else {
                v = parent[v];
            }
        }
        return v;
    }

    int lca(int v, int u) {
        if (depth[v] > depth[u]) {
            v = la(v, depth[v] - depth[u]);
        } else {
            u = la(u, depth[u] - depth[v]);
        }
        while (v != u) {
            if (up[v] != up[u]) {
                v = up[v];
                u = up[u];
            } else {
                v = parent[v];
                u = parent[u];
            }
        }
        return v;
    }
};

vector<vector<int>> tr, g, gg;
vector<int> tin, tout, c, ggg;
int timer = 0;
LCA st;
constexpr int LG = 18;

void add(int v, int p) {
    g.emplace_back();
    ggg[v] = g.size() - 1;
    g.back().push_back(c[v]);
    gg[0][v] = ggg[p];
    for (int i = 1; i < LG; ++i) {
        g.emplace_back();
        gg[i][v] = g.size() - 1;
        g.back().emplace_back(gg[i - 1][v]);
        g.back().emplace_back(gg[i - 1][st.la(v, 1 << (i - 1))]);
    }
}

void dfs(int v, int p) {
    tin[v] = timer++;
    if (p == -1) {
        g.emplace_back();
        ggg[v] = g.size() - 1;
        g.back().push_back(c[v]);
    }
    for (auto u : tr[v]) {
        if (u != p) {
            st.add_leaf(u, v);
            add(u, v);
            dfs(u, v);
        }
    }
    tout[v] = timer;
}

vector<int> val;
vector<int> curst;
vector<int> cond, siz;
int cnt = 0;
int k;

int dfs_scc(int j) {
    static int timer = 0;
    int low = val[j] = ++timer;
    curst.push_back(j);
    for (auto e : g[j]) {
        if (cond[e] < 0) {
            low = min(low, val[e] ? val[e] : dfs_scc(e));
        }
    }
    if (low == val[j]) {
        int x;
        do {
            x = curst.back();
            curst.pop_back();
            siz[cnt] += (x < k);
            cond[x] = cnt;
        } while (x != j);
        ++cnt;
    }
    return val[j] = low;
}

int main() {
    cin.tie(0)->sync_with_stdio(0);
#ifdef __APPLE__
    freopen("input.txt", "r", stdin);
#endif
    int n;
    cin >> n >> k;
    tr.resize(n);
    for (int i = 0; i < n - 1; ++i) {
        int u, v;
        cin >> u >> v;
        --u;
        --v;
        tr[u].push_back(v);
        tr[v].push_back(u);
    }
    c.resize(n);
    cin >> c;
    for (auto &i : c) {
        --i;
    }
    tin.resize(n);
    tout.resize(n);
    st = LCA(n);
    g.reserve(LG * n + k);
    g.resize(k);
    gg.resize(LG, vector<int>(n));
    ggg.resize(n);
    dfs(0, -1);
    vector<vector<int>> cc(k);
    for (int i = 0; i < n; ++i) {
        cc[c[i]].push_back(i);
    }
    auto add_edges = [&](int v, int d) {
        int cv = c[v];
        // for (int i = 0; i < d; ++i) {
        //     v = st.up[0][v];
        //     g[cv].push_back(c[v]);
        // }
        // if (cv == 2) {
        //     cout << "nigga\n";
        // }
        for (int i = LG - 1; i >= 0; --i) {
            if (d >= (1 << i)) {
                g[cv].push_back(gg[i][v]);
                d -= (1 << i);
                v = st.la(v, (1 << i));
            }
        }
    };
    for (int i = 0; i < k; ++i) {
        sort(cc[i].begin(), cc[i].end(), [&](int u, int v) {
            return tin[u] < tin[v];
        });
        for (int j = 0; j + 1 < (int) cc[i].size(); ++j) {
            int l = st.lca(cc[i][j], cc[i][j + 1]);
            add_edges(cc[i][j], st.depth[cc[i][j]] - st.depth[l]);
            add_edges(cc[i][j + 1], st.depth[cc[i][j + 1]] - st.depth[l]);
        }
    }
    // for (int i = 0; i < (int) g.size(); ++i) {
    //     for (int j : g[i]) {
    //         cout << i + 1 << " " << j + 1 << "\n";
    //     }
    // }
    gg.clear();
    gg.shrink_to_fit();
    int sz = g.size();
    val.resize(sz);
    cond.resize(sz, -1);
    siz.resize(sz);
    for (int i = 0; i < sz; ++i) {
        if (cond[i] < 0) {
            dfs_scc(i);
        }
    }
    vector<int> iznas(cnt);
    for (int i = 0; i < sz; ++i) {
        for (auto j : g[i]) {
            if (cond[i] != cond[j]) {
                ++iznas[cond[i]];
            }
        }
    }
    // debug(siz);
    // debug(iznas);
    int ans = INF;
    for (int i = 0; i < cnt; ++i) {
        if (iznas[i] == 0) {
            ans = min(ans, siz[i]);
        }
    }
    cout << ans - 1 << "\n";
}

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