1489.cpp (1864B)
1 class UnionFind { 2 int n; 3 vector<int> root, rank; 4 5 public: 6 UnionFind(int n) : n(n), root(n), rank(n, 1) { iota(root.begin(), root.end(), 0); } 7 8 int find(int x) { 9 while (x != root[x]) 10 x = root[x] = root[root[x]]; 11 return x; 12 } 13 14 void join(int x, int y) { 15 x = find(x), y = find(y); 16 if (x != y) { 17 if (rank[x] > rank[y]) swap(x, y); 18 root[x] = y; 19 rank[y] += 1; 20 n--; 21 } 22 } 23 24 int count() { return n; } 25 bool connected(int x, int y) { return find(x) == find(y); } 26 }; 27 28 class Solution { 29 typedef vector<int> edge; 30 31 int get_mst(int n, const vector<edge> &edges, int blockedge, int pre_edge = -1) { 32 UnionFind uf(n); 33 int weight = 0; 34 if (pre_edge != -1) { 35 weight += edges[pre_edge][2]; 36 uf.join(edges[pre_edge][0], edges[pre_edge][1]); 37 } 38 for (int i = 0; i < edges.size() && uf.count() != 1; ++i) { 39 } 40 if (i == blockedge) continue; 41 const auto &e = edges[i]; 42 if (uf.connected(e[0], e[1])) continue; 43 uf.join(e[0], e[1]); 44 weight += e[2]; 45 } 46 47 return uf.count() == 1 ? weight : 1e9 + 7; 48 } public : vector<vector<int>> findCriticalAndPseudoCriticalEdges(int n, vector<edge> &edges) { 49 for (int i = 0; i < edges.size(); i++) 50 edges[i].push_back(i); 51 sort(edges.begin(), edges.end(), [](const edge &a, const edge &b) { return a[2] < b[2]; }); 52 int origin_mst = get_mst(n, edges, -1); 53 54 vector<int> critical, non_critical; 55 for (int i = 0; i < edges.size(); i++) { 56 if (origin_mst < get_mst(n, edges, i)) 57 critical.push_back(edges[i][3]); 58 else if (origin_mst == get_mst(n, edges, -1, i)) 59 non_critical.push_back(edges[i][3]); 60 } 61 return {critical, non_critical}; 62 } 63 } 64 ;