Submission #677064

#	Time	Username	Problem	Language	Result	Execution time	Memory
677064		Cyanmond	Railway (BOI17_railway)	C++17	100 / 100	334 ms	65736 KiB

railway

#include <bits/stdc++.h>

struct LCA {
    int n;
    int logn;
    std::vector<std::vector<int>> tree;
    std::vector<int> depth;
    std::vector<int> parent;
    std::vector<std::vector<int>> ancestors;

    LCA(const std::vector<std::vector<int>> &treee) {
        tree = treee;
        n = (int)tree.size();
        logn = 0;
        while ((1 << logn) < n) {
            ++logn;
        }
        parent.resize(n);
        depth.resize(n);
        dfs(0, -1, 0);
        make_table();
    }

    void make_table() {
        ancestors.resize(logn + 1);
        ancestors[0] = parent;
        for (int d = 0; d < logn; ++d) {
            ancestors[d + 1].resize(n);
            for (int i = 0; i < n; ++i) {
                const int t = ancestors[d][i];
                if (t == -1) {
                    ancestors[d + 1][i] = -1;
                } else {
                    ancestors[d + 1][i] = ancestors[d][t];
                }
            }
        }
    }

    void dfs(int v, int p, int d) {
        parent[v] = p;
        depth[v] = d;
        for (const int t : tree[v]) {
            if (t == p) {
                continue;
            }
            dfs(t, v, d + 1);
        }
    }

    int lca(int a, int b) {
        int da = depth[a], db = depth[b];
        if (da > db) {
            std::swap(a, b);
            std::swap(da, db);
        }
        for (int d = logn; d >= 0; --d) {
            if (db - (1 << d) >= da) {
                b = ancestors[d][b];
                db -= (1 << d);
            }
        }
        assert(da == db);
        if (a == b) {
            return a;
        }
        for (int d = logn; d >= 0; --d) {
            int na = ancestors[d][a];
            int nb = ancestors[d][b];
            if (na != nb) {
                a = na;
                b = nb;
            }
        }
        return parent[a];
    }
};

void merge(std::set<int> &a, std::set<int> &b) {
    if (a.size() < b.size()) {
        std::swap(a, b);
    }
    for (const int e : b) {
        a.insert(e);
    }
}

int main() {
    int N, M, K;
    std::cin >> N >> M >> K;

    std::vector<int> A(N - 1), B(N - 1);
    std::vector<std::vector<int>> tree(N);
    std::map<std::pair<int, int>, int> edge_id;
    for (int i = 0; i < N - 1; ++i) {
        std::cin >> A[i] >> B[i];
        tree[--A[i]].push_back(--B[i]);
        tree[B[i]].push_back(A[i]);
        edge_id[{A[i], B[i]}] = edge_id[{B[i], A[i]}] = i;
    }

    std::vector<int> S(M);
    std::vector<std::vector<int>> L(M);
    for (int i = 0; i < M; ++i) {
        std::cin >> S[i];
        L[i].resize(S[i]);
        for (auto &e : L[i]) {
            std::cin >> e;
            --e;
        }
    }

    LCA tree_d(tree);
    std::vector<std::vector<int>> add_list(N), erase_list(N);
    for (int i = 0; i < M; ++i) {
        for (const int e : L[i]) {
            add_list[e].push_back(i);
        }

        int lca = L[i][0];
        for (int j = 1; j < S[i]; ++j) {
            lca = tree_d.lca(lca, L[i][j]);
        }
        erase_list[lca].push_back(i);
    }

    std::vector<int> answer;
    auto dfs = [&](auto &&self, const int v, const int p) -> std::set<int> {
        std::set<int> ls;
        for (const int t : tree[v]) {
            if (p == t) {
                continue;
            }
            auto res = self(self, t, v);
            if ((int)res.size() >= K) {
                answer.push_back(edge_id[{t, v}]);
            }
            merge(ls, res);
        }

        for (const int e : add_list[v]) {
            ls.insert(e);
        }

        for (const int e : erase_list[v]) {
            ls.erase(e);
        }

        return ls;
    };
    dfs(dfs, 0, -1);

    std::sort(answer.begin(), answer.end());
    std::cout << answer.size() << std::endl;
    for (int i = 0; i < (int)answer.size(); ++i) {
        std::cout << answer[i] + 1 << (i == (int)answer.size() - 1 ? '\n' : ' ');
    }
}

• Subtask 1: 8 / 8
• Subtask 2: 15 / 15
• Subtask 3: 7 / 7
• Subtask 4: 29 / 29
• Subtask 5: 16 / 16
• Subtask 6: 25 / 25