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

Link

MST variantions

TBA

class UF:
  def __init__(self, n):
    self.count = n
    self.parents = list(range(n))
  
  def find(self, x):
    if x != self.parents[x]:
      self.parents[x] = self.find(self.parents[x])

    return self.parents[x]

  def union(self, x, y):
    x = self.find(x)
    y = self.find(y)

    if x < y:
      self.parents[y] = x
    else:
      self.parents[x] = y
    
    self.count -= 1

class Solution:
    def findCriticalAndPseudoCriticalEdges(self, N: int, edges: List[List[int]]) -> List[List[int]]:

        E = len(edges)

        e = sorted(enumerate(edges), key=lambda x: x[1][2])

        def mst(idx, start):
            total = 0
            uf = UF(N)
            if start:
                uf.union(edges[idx][0], edges[idx][1])
                total += edges[idx][2]
            
            for i, [v1,v2,w] in e:
                if i == idx:
                    continue

                if uf.find(v1) != uf.find(v2):
                    uf.union(v1, v2)
                    total += w

            return total
            


        base_mst = mst(-1, False)

        critical = []
        pseudo = []
        for i in range(E):
          if mst(i, True) == base_mst:
            if mst(i, False) != base_mst:
              critical.append(i)
            else:
              pseudo.append(i)


        
        return [critical, pseudo]

        
        # M[v] = (w, i)