Submission #883103

#	Time	Username	Problem	Language	Result	Execution time	Memory
883103		vjudge1	Capital City (JOI20_capital_city)	C++17	11 / 100	2698 ms	524288 KiB

capital_city

// https://oj.uz/problem/view/JOI20_capital_city

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

namespace std {

template <class Fun>
class y_combinator_result {
        Fun fun_;

       public:
        template <class T>
        explicit y_combinator_result(T &&fun) : fun_(std::forward<T>(fun)) {}

        template <class... Args>
        decltype(auto) operator()(Args &&...args) {
                return fun_(std::ref(*this), std::forward<Args>(args)...);
        }
};

template <class Fun>
decltype(auto) y_combinator(Fun &&fun) {
        return y_combinator_result<std::decay_t<Fun>>(std::forward<Fun>(fun));
}

/* Example
        auto fun = y_combinator([&](auto self, int x) -> void {
                self(x + 1);
        });
*/

}  // namespace std

int32_t main() {
        ios_base::sync_with_stdio(0);
        cin.tie(0);

        int n, k;
        cin >> n >> k;
        vector<vector<int>> adj(n);
        for (int i = 0; i < n - 1; i++) {
                int u, v;
                cin >> u >> v;
                u--, v--;
                adj[u].emplace_back(v);
                adj[v].emplace_back(u);
        }
        vector<int> c(n);
        for (int i = 0; i < n; i++) cin >> c[i], c[i]--;

        vector<int> d(n);
        vector<vector<int>> up(__lg(n) + 1, vector<int>(n));
        vector<vector<int>> rmq(__lg(n) + 1, vector<int>(n));
        vector<vector<int>> g(k);
        vector<vector<int>> rg(k);

        auto dfs = y_combinator([&](auto self, int u, int p) -> void {
                up[0][u] = p;
                rmq[0][u] = c[u];
                for (int i = 1; i <= __lg(n); i++) {
                        up[i][u] = up[i - 1][up[i - 1][u]];
                        rmq[i][u] = g.size();
                        g.emplace_back(vector<int>({rmq[i - 1][u], rmq[i - 1][up[i - 1][u]]}));
                        rg.emplace_back();
                        rg[rmq[i - 1][u]].emplace_back(rmq[i][u]);
                        rg[rmq[i - 1][up[i - 1][u]]].emplace_back(rmq[i][u]);
                }
                for (int v : adj[u]) {
                        if (v == p) continue;
                        d[v] = d[u] + 1;
                        self(v, u);
                }
        });

        auto lca = [&](int group_id, int x, int y) {
                if (d[x] < d[y]) swap(x, y);
                for (int i = __lg(n); i >= 0; i--) {
                        if (d[up[i][x]] >= d[y]) {
                                g[group_id].emplace_back(rmq[i][x]);
                                rg[rmq[i][x]].emplace_back(group_id);
                                x = up[i][x];
                        }
                }
                if (x == y) {
                        g[group_id].emplace_back(rmq[0][x]);
                        rg[rmq[0][x]].emplace_back(group_id);
                        return x;
                }
                for (int i = __lg(n); i >= 0; i--) {
                        if (up[i][x] != up[i][y]) {
                                g[group_id].emplace_back(rmq[i][x]);
                                g[group_id].emplace_back(rmq[i][y]);
                                rg[rmq[i][x]].emplace_back(group_id);
                                rg[rmq[i][y]].emplace_back(group_id);
                                x = up[i][x], y = up[i][y];
                        }
                }
                g[group_id].emplace_back(rmq[0][x]);
                g[group_id].emplace_back(rmq[1][y]);
                rg[rmq[0][x]].emplace_back(group_id);
                rg[rmq[1][y]].emplace_back(group_id);
                return up[0][x];
        };

        vector<vector<int>> group(k);

        for (int i = 0; i < n; i++) group[c[i]].emplace_back(i);

        dfs(0, 0);

        for (int i = 0; i < k; i++) {
                int top = group[i][0];
                for (int j : group[i]) top = lca(i, top, j);
        }

        int N = g.size();
        vector<int> done(N), topo, comp(N);

        auto dfs2 = y_combinator([&](auto self, int u) -> void {
                done[u] = 1;
                for (int v : g[u]) {
                        if (!done[v]) self(v);
                }
                topo.emplace_back(u);
        });

        auto scc = y_combinator([&](auto self, int u, int id) -> void {
                done[u] = 0;
                comp[u] = id;
                for (int v : rg[u]) {
                        if (done[v]) self(v, id);
                }
        });

        for (int i = 0; i < N; i++) {
                if (!done[i]) dfs2(i);
        }

        reverse(topo.begin(), topo.end());

        int M = 0;
        for (int i : topo) {
                if (done[i]) scc(i, M++);
        }

        vector<int> deg(M), sz(M);
        vector<vector<int>> a(M);

        for (int i = 0; i < N; i++) {
                for (int j : g[i]) {
                        if (comp[i] == comp[j]) continue;
                        a[comp[i]].emplace_back(comp[j]);
                }
        }

        vector<int> vis(M);
        vector<int> cnt(M);
        for (int i = 0; i < k; i++) cnt[comp[i]] = 1, sz[comp[i]]++;

        auto dfs3 = y_combinator([&](auto self, int u) -> void {
                vis[u] = 1;
                for (int v : a[u]) {
                        if (!vis[v]) self(v);
                        cnt[u] += cnt[v];
                        cnt[u] = min(cnt[u], 2);
                }
        });

        for (int i = 0; i < M; i++) {
                if (!vis[i]) dfs3(i);
        }

        int res = k - 1;

        for (int i = 0; i < k; i++) {
                if (cnt[comp[i]] == 1) res = min(res, sz[comp[i]] - 1);
        }

        cout << res;
}

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