leetcode--找到最小生成树里的关键边和伪关键边

 题目是LeetCode第194场周赛的第四题，链接：1489. 找到最小生成树里的关键边和伪关键边。具体描述见原题。

 这题关键在于怎么找最小生成树，只要怎么找，就可以先生成一棵最小生成树同时记录这个最小权重和。然后只要把某条边去掉再重新执行生成最小生成树的过程，如果此生成的最小生成树的权重和大于原先的最小权重和，那么就说明这条边为关键边，否则则为伪关键边。生成最小生成树的算法的话可以直接用kruskal算法，具体就是先对所有边根据权重进行排序，然后每次都取出最小权重的边加入最小生成树中去，判断是否会形成环路（这个用并查集做），如果不形成环路就可以把这条边加入最小生成树，否则直接抛弃这条边（因为环路里其他边的权重都比当前这条边先加入，所以权重都比这条边小，自然需要去掉当前这条边）。假设节点数为n，边数为m，则时间复杂度为$O(m^2)$的（排序（mlogm）+m次kruskal+并查集（m次合并大概接近线性O(m)）），空间复杂度为$O(n+m)$。

 JAVA版代码如下：

class Solution {
    private int[] prev;
    private int[][] backup;

    private int find(int node) {
        if (prev[node] != prev[prev[node]]) {
            prev[node] = find(prev[node]);
        }
        return prev[node];
    }

    public int kruskal(int[][] edges, int n, int k, boolean removeK) {
        for (int i = 0; i < n; ++i) {
            prev[i] = i;
        }
        int sum = 0;
        int count = 0;
        if (!removeK && k >= 0) {
            prev[backup[k][1]] = backup[k][0];
            sum += backup[k][2];
            ++count;
        }
        for (int i = 0; i < edges.length; ++i) {
            if (removeK && edges[i] == backup[k]) {
                continue;
            }
            int prev1 = find(edges[i][0]);
            int prev2 = find(edges[i][1]);
            if (prev1 == prev2) {
                continue;
            }
            prev[prev2] = prev1;
            sum += edges[i][2];
            if (++count == n - 1) {
                return sum;
            }
        }
        return Integer.MAX_VALUE;
    }

    public List<List<Integer>> findCriticalAndPseudoCriticalEdges(int n, int[][] edges) {
        prev = new int[n];
        backup = edges.clone();
        Arrays.sort(edges, new Comparator<int[]>() {
            @Override
            public int compare(int[] i1, int[] i2) {
                return i1[2] - i2[2];
            }
        });
        int minSum = kruskal(edges, n, -1, false);
        List<List<Integer>> result = new LinkedList<>();
        result.add(new LinkedList<Integer>());
        result.add(new LinkedList<Integer>());
        for (int i = 0; i < edges.length; ++i) {
            if (kruskal(edges, n, i, true) > minSum) {
                result.get(0).add(i);
            }
            else if (kruskal(edges, n, i, false) == minSum) {
                result.get(1).add(i);
            }
        }
        return result;
    }
}

 提交结果如下:

 Python版代码如下：

class Solution:
    def __init__(self):
        self.MAX_VALUE = 1e6

    def find(self, prev, node):
        # 并查集的查找操作
        if prev[node] != node:
            # 路径压缩
            prev[node] = self.find(prev, prev[node])
        return prev[node]
    
    def kruskal(self, edges, n, k, remove=True):
        # 找最小生成树
        prev = [i for i in range(n)]
        count, cost = 0, 0
        if not remove:
            # 必包含第k条边
            prev[self.backup[k][1]] = self.backup[k][0]
            cost = self.backup[k][2]
            count = 1
        
        for i, edge in enumerate(edges):
            if remove and k >= 0 and edge == self.backup[k]:
                # 不包括第k条边
                continue
            
            p1, p2 = self.find(prev, edge[0]), self.find(prev, edge[1])
            if p1 == p2:
                # 如果位于同一棵树，避免成环直接跳过
                continue
            # 两棵树的合并
            prev[p2] = p1
            cost += edge[2]
            count += 1
            if count == n - 1:
                return cost
        return self.MAX_VALUE

    def findCriticalAndPseudoCriticalEdges(self, n: int, edges: List[List[int]]) -> List[List[int]]:
        # 只是保持顺序，不用深拷贝
        self.backup = [edge for edge in edges]
        # 按权值升序排序
        edges.sort(key=lambda x : x[2])
        minCost = self.kruskal(edges, n, -1)

        result = [[], []]
        for i in range(len(edges)):
            if self.kruskal(edges, n, i) > minCost:
                # 关键边
                result[0].append(i)
            elif self.kruskal(edges, n, i, False) == minCost:
                # 伪关键边
                result[1].append(i)
        return result

 提交结果如下: