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

Given a weighted undirected connected graph with n vertices numbered from 0 to n - 1, and an array edges where edges[i] = [ai, bi, weighti] represents a bidirectional and weighted edge between nodes ai and bi. A minimum spanning tree (MST) is a subset of the graph’s edges that connects all vertices without cycles and with the minimum possible total edge weight.

Find all the critical and pseudo-critical edges in the given graph’s minimum spanning tree (MST). An MST edge whose deletion from the graph would cause the MST weight to increase is called a critical edge. On the other hand, a pseudo-critical edge is that which can appear in some MSTs but not all.

Note that you can return the indices of the edges in any order.

class Solution {
    int find(vector<int>& g, int n) {
        return g[n] == n ? n : g[n] = find(g,g[n]);
    }
    void uni(vector<int>& g, int u, int v) {
        int pu = find(g,u), pv = find(g,v);
        g[pu] = g[pv] = min(pu, pv);
    }
    int mstCost(int n, vector<vector<int>>& edges, int blockEdge, int reqEdge) {
        int cost = 0;
        vector<int> g(n);
        for(int i = 0; i < n; i++) g[i] = i;

        if(reqEdge != -1) {
            uni(g, edges[reqEdge][0], edges[reqEdge][1]);
            cost += edges[reqEdge][2];
        }
        
        for(int i = 0; i < edges.size(); i++) {
            if(i == blockEdge or find(g, edges[i][0]) == find(g, edges[i][1])) continue;
            uni(g, edges[i][0], edges[i][1]);
            cost += edges[i][2];
        }
        
        for(int i = 0; i < n; i++) {
            if(find(g,i)) return 1e9;
        }

        return cost;
    }
public:
    vector<vector<int>> findCriticalAndPseudoCriticalEdges(int n, vector<vector<int>>& edges) {
        for(int i = 0; i < edges.size(); i++) edges[i].push_back(i);
        sort(edges.begin(), edges.end(), [](vector<int>& e1, vector<int>& e2) {
            return e1[2] < e2[2];
        });
        int cost = mstCost(n, edges, -1, -1);
        
        vector<int> critical, nonCritical;

        for(int i = 0; i < edges.size(); i++) {
            if(cost < mstCost(n, edges, i, -1)) {
                critical.push_back(edges[i][3]);
            } else if(cost == mstCost(n, edges, -1, i)) {
                nonCritical.push_back(edges[i][3]);
            }
        }

        return {critical, nonCritical};
    }
};