[Python] UnionFind + Kruskal solution with detail explanation - from O(E^2) to O(ElogE)
Background\n\nThe Kruskal\u2019s algorithm finds the minimum cost spanning tree using a greedy approach: it first sorts all edges, which takes O(ELogE) time. Af
**Background**\n\nThe Kruskal\u2019s algorithm finds the minimum cost spanning tree using a greedy approach: it first sorts all edges, which takes `O(ELogE)` time. After sorting, it iterates through all edges and apply Union-Find algorithm. \n\nThe Union-Find by Rank and Path Compression  takes `O(alpha(V) * E)`, where `alpha(V)` is the inverse of the Ackermann function and can be considered constant.\n\n<br />\n\n**Explaination - O(E^2) Solution**\n\nFrom the Background section, we know that the bottleneck for Kruskal\'s algorithm is really the sorting part. If all edges are already sorted, we can find the MST with any subset of nodes in `O(E)` time.\n\nWe can take the following steps to solve this problem:\n\n1. Sort all the edges in non-decreasing order of their weight.\n2. Run Kruskal\u2019s algorithm to find the total weight of the MST (i.e. `base`).\n3. Iterate through each edge and find the weight of MST if we excludes this edge (i.e. `wgt_excl`).\n4. If `wgt_excl != base`, this edge is a critical edge.\n5. If `wgt_excl == base`, this edge is either pseudo-critical (MST may or may not include this edge), or redundant (MST cannot include this edge). Let\'s find the MST\'s total weights if we do include this edge (i.e. `wgt_incl`). If `wgt_incl` is larger than `base`, then this edge is redundant; otherwise, this edge is pseudo-critial.\n\nThe time complexity is `O(E^2)`, because we do sorting once (`O(ElogE)`) and calculating MST `2E` times (`O(E^2)`) .\n\n<br />\n\n**Explaination - O(ElogE) Solution**\n\nIn the previous algorithm, we calculated MST for `2E` times. This takes a lot of time and is not necessary. Instead, we can populate the list of critical edges and pesudo-critical edges when we calculates the MST. In this way, we only need to calculate the MST once.\n\nI will list the general thoughts below, but please see code comments for more details.\n\nFirst of all, we sort all edges in the same way as we did in the Kruskal\u2019s algorithm.\n\nThen, for a given weight, we check all its corresponding edges. For edge `(u, v)`, if it connects two already connected groups, it is redundant (neither critical nor pesudo-critical); otherwise, we put `(u, v)`\'s group parents `(pu, pv)` in `seen`.\n\nIf there is more than one edge that connects union group `pu` and `pv`, then all `seen[pu, pv]` are pesudo-critical. If not, we check if the edges in `seen[pu, pv]` forms any cycles. For the edges not forming cycles, we mark them as critical, and mark the others as pesudo-critical.\n\nTo identify cycles in a graph, we can refer to the solution of [1192. Critical Connections in a Network](https://leetcode.com/problems/critical-connections-in-a-network/).\n\n\n<br />\n\n**O(E^2) Solution**\n```Python\nclass UnionFindSet:\n    def __init__(self, n=0):\n        self.parents = {}\n        self.ranks = {}\n        self.count = 0\n        for i in range(n):\n            self.add(i)\n\n    def add(self, p):\n        self.parents[p] = p\n        self.ranks[p] = 1\n        self.count += 1\n\n    def find(self, u):\n        if u != self.parents[u]:\n            self.parents[u] = self.find(self.parents[u])\n        return self.parents[u]\n\n    def union(self, u, v):\n        pu, pv = self.find(u), self.find(v)\n        if pu == pv: \n            return False\n        if self.ranks[pu] < self.ranks[pv]:\n            self.parents[pu] = pv\n        elif self.ranks[pu] > self.ranks[pv]:\n            self.parents[pv] = pu\n        else:        \n            self.parents[pv] = pu\n            self.ranks[pu] += 1\n        self.count -= 1\n        return True\n    \n# UnionFind + Kruskal + Enumerate edges\n# O(ElogE + E^2 + E^2)\nclass Solution:\n    def findCriticalAndPseudoCriticalEdges(self, n: int, edges: List[List[int]]) -> List[List[int]]:\n        # sort edges in asc order based on weight\n        edges = [(u, v, w, i) for i, (u, v, w) in enumerate(edges)]\n        edges.sort(key=lambda x: x[2])\n        \n        # do not use this edge\n        def find_mst_without_this_edge(edge_idx):\n            union_find_set = UnionFindSet(n)\n            ans = 0\n            for i, (u, v, w, _) in enumerate(edges):\n                # do not use this edge\n                if i == edge_idx:\n                    continue\n                if union_find_set.union(u, v):\n                    ans += w\n            parent = union_find_set.find(0)\n            return ans if all(union_find_set.find(i) == parent for i in range(n)) else inf\n        \n        # need to use this edge\n        def find_mst_with_this_edge(edge_idx):\n            union_find_set = UnionFindSet(n)\n            # use this edge first\n            u0, v0, w0, _ = edges[edge_idx]\n            ans = w0\n            union_find_set.union(u0, v0)\n            for i, (u, v, w, _) in enumerate(edges):\n                # do not use this edge\n                if i == edge_idx:\n                    continue\n                if union_find_set.union(u, v):\n                    ans += w\n            parent = union_find_set.find(0)\n            return ans if all(union_find_set.find(i) == parent for i in range(n)) else inf\n        \n        # normal MST total weight\n        base = find_mst_without_this_edge(-1)\n        cri, p_cri = set(), set()\n        for i in range(len(edges)):\n            wgt_excl = find_mst_without_this_edge(i)\n            # if not included, MST total weight would increase\n            if wgt_excl > base:\n                cri.add(edges[i][3])\n            else:\n                wgt_incl = find_mst_with_this_edge(i)\n                # with this edge, MST total weight doesn\'t change\n                if wgt_incl == base:\n                    p_cri.add(edges[i][3])\n    \n        return [cri, p_cri]\n```\n\n\n\n**O(ElogE) Solution**\n```Python\nclass Solution:\n    def findCriticalAndPseudoCriticalEdges(self, n: int, edges: List[List[int]]) -> List[List[int]]:\n    \n        # reference: LC 1192\n        def dfs(curr, level, parent):\n            levels[curr] = level\n            for child, i in graph[curr]:\n                if child == parent:\n                    continue\n                elif levels[child] == -1:\n                    levels[curr] = min(levels[curr], dfs(child, level + 1, curr))\n                else:\n                    levels[curr] = min(levels[curr], levels[child])\n                if levels[child] >= level + 1 and i not in p_cri:\n                    cri.add(i)\n            return levels[curr]\n        \n        # init critical and pseudo-critical edge set\n        cri, p_cri = set(), set()\n        \n        # use dic to store all edges associated with a given weight\n        dic = collections.defaultdict(list)\n        for i, (u, v, w) in enumerate(edges):\n            dic[w].append([u, v, i])\n        \n        # define union find et\n        union_set = UnionFindSet(n)\n        \n        # iterate through all weights in ascending order\n        for w in sorted(dic):\n                \n            # seen[(pu, pv)] contains all edges connecting pu and pv,\n            # where pu and pv are the parent nodes of their corresponding groups\n            seen = collections.defaultdict(set)\n            # populate seen\n            for u, v, i in dic[w]:\n                pu, pv = union_set.find(u), union_set.find(v)\n                # skip the edge that creates cycle\n                if pu == pv:\n                    continue\n                seen[min(pu, pv), max(pu, pv)].add(i) # edge i connects pu and pv\n            \n            # w_edges contains all weight-w edges we may add to MST\n            w_edges, graph = [], collections.defaultdict(list)\n            for pu, pv in seen:\n                # more than 1 edge can connect pu and pv\n                # these edges are pseudo-critical\n                if len(seen[pu, pv]) > 1:\n                    p_cri |= seen[pu, pv]\n                # construct graph for weight w \n                edge_idx = seen[pu, pv].pop()\n                graph[pu].append((pv, edge_idx))\n                graph[pv].append((pu, edge_idx))\n                w_edges.append((pu, pv, edge_idx))\n                # union pu and pv groups\n                union_set.union(pu, pv)\n            \n            # run dfs to mark all critical w_edges\n            levels = [-1] * n\n            for u, v, i in w_edges:\n                if levels[u] == -1:\n                    dfs(u, 0, -1)\n            # the edges in w_edges cycles are pseudo-critical\n            for u, v, i in w_edges:\n                if i not in cri:\n                    p_cri.add(i)\n        \n        return [cri, p_cri]\n```
[]