Find Critical and Pseudo-Critical Edges in Minimum Spanning Tree LeetCode solution in Python

1
class Solution:
2
    def findCriticalAndPseudoCriticalEdges(
3
        self, n: int, edges: List[List[int]]
4
    ) -> List[List[int]]:
5
        def find(v, parent):
6
            if parent[v] != v:
7
                parent[v] = find(parent[v], parent)
8
            return parent[v]
9

10
        def union(u, v, parent):
11
            parent[find(u, parent)] = find(v, parent)
12

13
        edges = [(u, v, w, i) for i, (u, v, w) in enumerate(edges)]
14
        edges.sort(key=lambda e: e[2])
15

16
        def find_mst_without_this_edge(idx):
17
            parent = list(range(n))
18
            res = 0
19
            for i, (u, v, w, _) in enumerate(edges):
20
                if i == idx:
21
                    continue
22
                if find(u, parent) != find(v, parent):
23
                    res += w
24
                    union(u, v, parent)
25
            root = find(0, parent)
26
            return (
27
                res if all(find(i, parent) == root for i in range(n)) else float("inf")
28
            )
29

30
        def find_mst_with_this_edge(idx):
31
            parent = list(range(n))
32
            u0, v0, w0, _ = edges[idx]
33
            res = w0
34
            union(u0, v0, parent)
35
            for i, (u, v, w, _) in enumerate(edges):
36
                if i == idx:
37
                    continue
38
                if find(u, parent) != find(v, parent):
39
                    res += w
40
                    union(u, v, parent)
41
            root = find(0, parent)
42
            return (
43
                res if all(find(i, parent) == root for i in range(n)) else float("inf")
44
            )
45

46
        base_mst_wgt = find_mst_without_this_edge(-1)
47
        cri, pcri = set(), set()
48
        for i in range(len(edges)):
49
            wgt_excl = find_mst_without_this_edge(i)
50
            if wgt_excl > base_mst_wgt:
51
                cri.add(edges[i][3])
52
            else:
53
                wgt_incl = find_mst_with_this_edge(i)
54
                if wgt_incl == base_mst_wgt:
55
                    pcri.add(edges[i][3])
56
        return [cri, pcri]