MST边集检查(MST Edge Set Check) 题解

思路分析

按题目要求使用对应数据结构与算法模拟/递推求解，注意数据范围边界与时间复杂度。

参考代码

#include <iostream>
#include <vector>
#include <numeric>
#include <algorithm>
using namespace std;

struct Edge {
    int u, v, w;
};

struct QGroup {
    int w, qid;
    vector<int> eids;
    bool operator<(const QGroup& o) const {
        if (w != o.w) return w < o.w;
        return qid < o.qid;
    }
};

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    int n, m, q;
    if (!(cin >> n >> m >> q)) return 0;
    vector<Edge> edges(m + 1);
    vector<int> ord(m);
    for (int i = 1; i <= m; ++i) {
        cin >> edges[i].u >> edges[i].v >> edges[i].w;
        ord[i - 1] = i;
    }
    sort(ord.begin(), ord.end(), [&](int a, int b) {
        return edges[a].w < edges[b].w;
    });
    vector<QGroup> groups;
    for (int i = 0; i < q; ++i) {
        int k;
        cin >> k;
        vector<int> list(k);
        for (int j = 0; j < k; ++j) cin >> list[j];
        sort(list.begin(), list.end(), [&](int a, int b) {
            return edges[a].w < edges[b].w;
        });
        for (int l = 0; l < k; ) {
            int r = l;
            while (r < k && edges[list[r]].w == edges[list[l]].w) ++r;
            QGroup g;
            g.w = edges[list[l]].w;
            g.qid = i;
            for (int t = l; t < r; ++t) g.eids.push_back(list[t]);
            groups.push_back(move(g));
            l = r;
        }
    }
    sort(groups.begin(), groups.end());
    vector<int> p(n + 1), sz(n + 1, 1);
    iota(p.begin(), p.end(), 0);
    auto find = [&](int x) {
        while (p[x] != x) x = p[x];
        return x;
    };
    vector<char> ok(q, 1);
    int ePtr = 0;
    int gSz = (int)groups.size();
    for (int i = 0; i < gSz; ) {
        int cw = groups[i].w;
        while (ePtr < m && edges[ord[ePtr]].w < cw) {
            int eid = ord[ePtr++];
            int ru = find(edges[eid].u), rv = find(edges[eid].v);
            if (ru != rv) {
                if (sz[ru] > sz[rv]) swap(ru, rv);
                p[ru] = rv;
                sz[rv] += sz[ru];
            }
        }
        int j = i;
        while (j < gSz && groups[j].w == cw) {
            if (ok[groups[j].qid]) {
                vector<pair<int, int>> rb;
                for (int eid : groups[j].eids) {
                    int ru = find(edges[eid].u), rv = find(edges[eid].v);
                    if (ru == rv) {
                        ok[groups[j].qid] = 0;
                        break;
                    }
                    if (sz[ru] > sz[rv]) swap(ru, rv);
                    p[ru] = rv;
                    sz[rv] += sz[ru];
                    rb.push_back({ru, rv});
                }
                while (!rb.empty()) {
                    auto [ru, rv] = rb.back();
                    rb.pop_back();
                    p[ru] = ru;
                    sz[rv] -= sz[ru];
                }
            }
            ++j;
        }
        while (ePtr < m && edges[ord[ePtr]].w == cw) {
            int eid = ord[ePtr++];
            int ru = find(edges[eid].u), rv = find(edges[eid].v);
            if (ru != rv) {
                if (sz[ru] > sz[rv]) swap(ru, rv);
                p[ru] = rv;
                sz[rv] += sz[ru];
            }
        }
        i = j;
    }
    for (int i = 0; i < q; ++i) cout << (ok[i] ? "YES\n" : "NO\n");
    return 0;
}