Submission #805385

#	Time	Username	Problem	Language	Result	Execution time	Memory
805385		FedShat	Capital City (JOI20_capital_city)	C++17	11 / 100	1859 ms	524288 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;
}

struct LCA {
    static constexpr int k = 18;
    vector<vector<int>> up;
    vector<int> depth;

    LCA() {}
    LCA(int n) {
        up.resize(k, vector<int>(n));
        depth.resize(n);
    }

    void add_leaf(int v, int p) {
        depth[v] = depth[p] + 1;
        up[0][v] = p;
        for (int i = 1; i < k; ++i) {
            up[i][v] = up[i - 1][up[i - 1][v]];
        }
    }

    int up_d(int v, int d) {
        for (int i = k - 1; i >= 0; --i) {
            if (d >= (1 << i)) {
                v = up[i][v];
                d -= (1 << i);
            }
        }
        return v;
    }

    int lca(int u, int v) {
        if (depth[u] < depth[v]) {
            swap(u, v);
        }
        u = up_d(u, depth[u] - depth[v]);
        for (int i = k - 1; i >= 0; --i) {
            if (up[i][u] != up[i][v]) {
                u = up[i][u];
                v = up[i][v];
            }
        }
        if (u == v) {
            return u;
        }
        return up[0][u];
    }
};

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

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 < LCA::k; ++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.up[i - 1][v]]);
    }
}

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<bool> used;
vector<int> topsort;

void dfs1(int v) {
    used[v] = 1;
    for (auto u : g[v]) {
        if (!used[u]) {
            dfs1(u);
        }
    }
    topsort.push_back(v);
}

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

void dfs2(int v) {
    used[v] = 1;
    cond[v] = cnt;
    siz[cnt] += (v < k);
    for (auto u : gr[v]) {
        if (!used[u]) {
            dfs2(u);
        }
    }
}

vector<vector<int>> gcond;

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.resize(k);
    gg.resize(LCA::k, 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 = LCA::k - 1; i >= 0; --i) {
            if (d >= (1 << i)) {
                g[cv].push_back(gg[i][v]);
                d -= (1 << i);
                v = st.up[i][v];
            }
        }
    };
    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";
    //     }
    // }
    int sz = g.size();
    gr.resize(sz);
    for (int i = 0; i < sz; ++i) {
        for (int j : g[i]) {
            gr[j].push_back(i);
        }
    }
    used.resize(sz);
    for (int i = 0; i < sz; ++i) {
        if (!used[i]) {
            dfs1(i);
        }
    }
    used.assign(sz, 0);
    cond.resize(sz);
    reverse(topsort.begin(), topsort.end());
    for (int i : topsort) {
        if (!used[i]) {
            siz.emplace_back();
            dfs2(i);
            ++cnt;
        }
    }
    gcond.resize(cnt);
    for (int i = 0; i < sz; ++i) {
        for (int j : g[i]) {
            if (cond[i] != cond[j]) {
                gcond[cond[i]].push_back(cond[j]);
            }
        }
    }
    for (int i = 0; i < cnt; ++i) {
        sort(gcond[i].begin(), gcond[i].end());
        gcond[i].resize(unique(gcond[i].begin(), gcond[i].end()) - gcond[i].begin());
    }
    used.assign(cnt, 0);
    int ans = INF;
    for (int i = 0; i < cnt; ++i) {
        if (gcond[i].empty()) {
            ans = min(ans, siz[i]);
        }
    }
    cout << ans - 1 << "\n";
}

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