Create 1489-find-critical-and-pseudo-critical-edges-in-minimum-spanning-tree.py

Author: neetcode-gh

python/1489-find-critical-and-pseudo-critical-edges-in-minimum-spanning-tree.py

@@ -0,0 +1,59 @@
+class UnionFind:
+    def __init__(self, n):
+        self.par = [i for i in range(n)]
+        self.rank = [1] * n
+
+    def find(self, v1):
+        while v1 != self.par[v1]:
+            self.par[v1] = self.par[self.par[v1]]
+            v1 = self.par[v1]
+        return v1
+
+    def union(self, v1, v2):
+        p1, p2 = self.find(v1), self.find(v2)
+        if p1 == p2:
+            return False
+        if self.rank[p1] > self.rank[p2]:
+            self.par[p2] = p1
+            self.rank[p1] += self.rank[p2]
+        else:
+            self.par[p1] = p2
+            self.rank[p2] += self.rank[p1]
+        return True
+
+class Solution:
+    def findCriticalAndPseudoCriticalEdges(self, n: int, edges: List[List[int]]) -> List[List[int]]:
+        # Time: O(E^2) - UF operations are assumed to be approx O(1)
+        for i, e in enumerate(edges):
+            e.append(i) # [v1, v2, weight, original_index]
+
+        edges.sort(key=lambda e: e[2])
+
+        mst_weight = 0
+        uf = UnionFind(n)
+        for v1, v2, w, i in edges:
+            if uf.union(v1, v2):
+                mst_weight += w
+
+        critical, pseudo = [], []
+        for n1, n2, e_weight, i in edges:
+            # Try without curr edge
+            weight = 0
+            uf = UnionFind(n)
+            for v1, v2, w, j in edges:
+                if i != j and uf.union(v1, v2):
+                    weight += w
+            if max(uf.rank) != n or weight > mst_weight:
+                critical.append(i)
+                continue
+            
+            # Try with curr edge
+            uf = UnionFind(n)
+            uf.union(n1, n2)
+            weight = e_weight
+            for v1, v2, w, j in edges:
+                if uf.union(v1, v2):
+                    weight += w
+            if weight == mst_weight:
+                pseudo.append(i)
+        return [critical, pseudo]