1489. 找到最小生成树里的关键边和伪关键边 Kruskal

力扣

class Solution {
public:
    vector<int> initDSU(int n)
    {
        //初始化并查集
        vector<int> ret;
        for (int i = 0; i < n; ++i) {
            ret.push_back(i);
        }
        return ret;
    }

    int find(int x, vector<int> &f)
    {
        if (x != f[x]) {
            f[x] = find(f[x], f);
        }
        return f[x];
    }

    void merge(int x, int y, vector<int> &f) 
    {
        int fx = find(x, f), fy = find(y, f);
        if (fx != fy) {
            f[fx] = fy;
        }
    }

    int kruskal(int n, int v, vector<vector<int> > &edges, vector<int> &f)
    {   
        for (auto e : edges) {
            if (e[0] < 0) 
                continue;   //被去除的边
            int fa = find(e[0], f), fb = find(e[1], f);
            if (fa != fb) {
                v += e[2];
                merge(e[0], e[1], f);
            }
        }

        for (int i = 0; i < n; ++i) {
            if (find(i, f) != f[0])
                return -1;
        }

        return v;
    }

    vector<vector<int>> findCriticalAndPseudoCriticalEdges(int n, vector<vector<int>> edges) 
    {   
        for (int i = 0; i < edges.size(); ++i) {
            edges[i].push_back(i);  //记录排序前的序号
        }
        sort(edges.begin(), edges.end(), [&](auto a, auto b){
            return a[2] < b[2];
        });

        auto f1 = initDSU(n);
        int value = kruskal(n, 0, edges, f1);
        vector<int> criticalEdges;
        vector<int> pseudoCriticalEdges;

        for (int i = 0; i < edges.size(); ++i) {
            edges[i].push_back(i);  //记录排序前的序号
        }

        sort(edges.begin(), edges.end(), [&](auto a, auto b){
            return a[2] < b[2];
        });

        for (int i = 0; i < edges.size(); ++i) {
            //判断是否是关键边
            auto f2 = initDSU(n);
            edges[i][0] -= 0x3f3f3f3f;   //去除这条边
            int v1 = kruskal(n, 0, edges, f2);
            edges[i][0] += 0x3f3f3f3f;
        
            if (v1 != value) {
                criticalEdges.push_back(edges[i][3]);
                continue;   //是关键边不需要再判断伪关键边
            }
            
            //判断是否是伪关键边
            auto f3 = initDSU(n);
            merge(edges[i][0], edges[i][1], f3);
            int v2 = kruskal(n, edges[i][2], edges, f3);
            if (v2 == value)
                pseudoCriticalEdges.push_back(edges[i][3]);
            
        }

        return {criticalEdges, pseudoCriticalEdges};
    }
};