Problem statement

https://binarysearch.com/problems/Minimum-Spanning-Tree/

Solution

It is special case of Leetcode 1489. Find Critical and Pseudo-Critical Edges in Minimum Spanning - the idea is to build two MST: one for original graph and another for graph with taken edge.

Complexity

Time complexity is O(E log E), space is O(E).

Code

class DSU:
    def __init__(self, N):
        self.p = list(range(N))

    def find(self, x):
        if self.p[x] != x:
            self.p[x] = self.find(self.p[x])
        return self.p[x]

    def union(self, x, y):
        xr = self.find(x)
        yr = self.find(y)
        self.p[xr] = yr

class Solution:
    def solve(self, edges, a, b):
        n = max(max(i for i, _, _ in edges), max(j for _, j, _ in edges)) + 1

        def kruskal(n, edges, taken):
            dsu, ans = DSU(n), 0
            for x, y, w in taken: 
                dsu.union(x, y)
                ans += w
            for x, y, w in sorted(edges, key = lambda x:x[2]):
                if dsu.find(x) == dsu.find(y): continue
                dsu.union(x, y)
                ans += w
            return ans

        MST1 = kruskal(n, edges, [])
        E2 = [(x, y, w) for x, y, w in edges if (x, y) not in [(a, b), (b, a)]]
        for x, y, w in edges:
            if (x, y) in [(a, b), (b, a)]: taken = [(x, y, w)]
        MST2 = kruskal(n, E2, taken)
        return MST1 == MST2