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

给你一个 n 个点的带权无向连通图，节点编号为 0 到 n-1 ，同时还有一个数组 edges ，其中 edges[i] = [fromi, toi, weighti] 表示在 fromi 和 toi 节点之间有一条带权无向边。最小生成树 (MST) 是给定图中边的一个子集，它连接了所有节点且没有环，而且这些边的权值和最小。

请你找到给定图中最小生成树的所有关键边和伪关键边。如果从图中删去某条边，会导致最小生成树的权值和增加，那么我们就说它是一条关键边。伪关键边则是可能会出现在某些最小生成树中但不会出现在所有最小生成树中的边。

请注意，你可以分别以任意顺序返回关键边的下标和伪关键边的下标。

来源：力扣（LeetCode）
链接：https://leetcode-cn.com/problems/find-critical-and-pseudo-critical-edges-in-minimum-spanning-tree
著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。

class Solution {
public:
    unordered_map <int, int> fa;
    unordered_map <int, int> rank;
    int count;

    int find(int x) {
        if (!fa.count(x)) {
            fa[x] = x;
            rank[x] = 1;
        }
        return fa[x] == x ? fa[x] : fa[x] = find(fa[x]);
    }

    bool same(int x, int y) {
        return find(x) == find(y);
    }

    bool unite(int x, int y) {
        int xx = find(x);
        int yy = find(y);
        if (same(xx, yy)) {
            return false;
        }
        if (rank[xx] < rank[yy]) {
            swap(xx, yy);
        }
        rank[xx] += rank[yy];
        fa[yy] = xx;
        count--;
        return true;
    }

    vector<vector<int>> findCriticalAndPseudoCriticalEdges(int n, vector<vector<int>>& edges) {
        vector <vector<int>> ans(2);

        int m = edges.size();
        // 记录每条边的id，后面输出关键边用
        for (int i = 0; i < m; i++) {
            edges[i].emplace_back(i);
        }

        // 正常kruskal求最小生成树
        int value = 0;
        sort(edges.begin(), edges.end(), [](const auto& u, const auto& v) {
            return u[2] < v[2];
        });
        for (int i = 0; i < m; i++) {
            if (unite(edges[i][0], edges[i][1])) {
                value += edges[i][2];
            }
        }

        // 判断是否是关键边
        for (int i = 0; i < m; i++) {
            // 初始化
            fa.clear();
            rank.clear();
            count = n;
            int v = 0;
            // kruskal
            for (int j = 0; j < m; j++) {
                if (j != i && unite(edges[j][0], edges[j][1])) {
                    v += edges[j][2];
                }
            }
            if (count != 1 || (count == 1 && v > value)) {
                ans[0].emplace_back(edges[i][3]);
                continue;
            }

            // 判断是否是伪关键边
            // 初始化
            fa.clear();
            rank.clear();
            count = n;
            v = 0;
            unite(edges[i][0], edges[i][1]);
            v = edges[i][2];

            for (int j = 0; j < m; j++) {
                if (j != i && unite(edges[j][0], edges[j][1])) {
                    v += edges[j][2];
                }
            }
            if (v == value) {
                ans[1].emplace_back(edges[i][3]);
            }
        }

        return ans;
    }
};