【Leetcode】1489. Find Critical and Pseudo-Critical Edges in Minimum Spanning Tree

题目地址：

https://leetcode.com/problems/find-critical-and-pseudo-critical-edges-in-minimum-spanning-tree/

给定一个无向带权图的所有边，求两个数组$A$和$B$，$A$存的是这样的边的下标：其在每一个这个图的最小生成树中都一定存在；$B$存的是这样的边的下标：其存在于某个最小生成树，但也存在某个最小生成树不含它。题目保证不含平行边，并且保证最小生成树是存在的。

思路是Kruskal算法。先求一下最小生成树的总边权，然后枚举每条边，先看一下如果去掉这条边，会不会导致最小生成树不存在或者导致最小生成树的总边权变大，如果是，则这条边是必须加的，所以将其下标加入$A$；如果不是，则需要看一下它是否要加入$B$，其实就是求是否存在含这条边的最小生成树，那么就先将这条边加入，然后看一下是否也能得到最小生成树，如果是，则说明存在某个最小生成树是含这条边的，则加入$B$。代码如下：

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

public class Solution {
    
    class UnionFind {
        int[] p;
        // group存的是当前并查集里有多少个不相交集
        int group;
        
        public UnionFind(int size) {
            p = new int[size];
            for (int i = 0; i < size; i++) {
                p[i] = i;
            }
            group = size;
        }
        
        public int find(int x) {
            if (p[x] != x) {
                p[x] = find(p[x]);
            }
            
            return p[x];
        }
        
        public void union(int x, int y) {
            int px = find(x), py = find(y);
            if (px != py) {
                p[px] = py;
                group--;
            }
        }
        
        public int getGroup() {
            return group;
        }
    }
    
    public List<List<Integer>> findCriticalAndPseudoCriticalEdges(int n, int[][] edges) {
        for (int i = 0; i < edges.length; i++) {
            int x = edges[i][0], y = edges[i][1], len = edges[i][2];
            edges[i] = new int[]{x, y, len, i};
        }
        
        List<List<Integer>> res = new ArrayList<>();
        res.add(new ArrayList<>());
        res.add(new ArrayList<>());
        
        // 将边按照边权从小到大排序
        Arrays.sort(edges, (e1, e2) -> Integer.compare(e1[2], e2[2]));
        // 求一下最小生成树
        int minCost = MST(edges, -1, -1, n);
        for (int i = 0; i < edges.length; i++) {
        	// 求如果去掉当前边得到的最小生成树总边权是多少
            int cost = MST(edges, i, -1, n), idx = edges[i][3];
            // 如果得到的总边权大于不删当前边的最小生成树总边权，则说明当前边必须加；
            // 否则看一下如果必须加入当前边，得到的最小生成树总边权是否等于没有限制下
            // 的最小生成树总边权，如果是，则说明存在含当前边的最小生成树
            if (cost > minCost) {
                res.get(0).add(idx);
            } else if (MST(edges, -1, i, n) == minCost) {
				res.get(1).add(idx);
            }
        }
        
        return res;
    }
    
    // deleted存的是不可用的边的下标，required存的是必须加的边的下标
    private int MST(int[][] edges, int deleted, int required, int n) {
        UnionFind uf = new UnionFind(n);
        int res = 0;
        if (required != -1) {
            int[] edge = edges[required];
            uf.union(edge[0], edge[1]);
            res += edge[2];
        }
        
        for (int i = 0; i < edges.length; i++) {
            if (i == required || i == deleted) {
                continue;
            }
            
            int[] edge = edges[i];
            int x = edge[0], y = edge[1];
            if (uf.find(x) != uf.find(y)) {
                uf.union(x, y);
                res += edge[2];
                
                if (uf.getGroup() == 1) {
                    return res;
                }
            }
        }
        
        // 如果不存在最小生成树，则返回正无穷
        return Integer.MAX_VALUE;
    }
}

时间复杂度$O(m^2{\log }^{\ast }n)$，$n$和$m$分别是图的顶点数和边数，空间复杂度$O(n+m)$。