leetcode

0 class UnionFind {

1 int n;

2 vector<int> root, rank;

3

4 public:

5 UnionFind(int n) : n(n), root(n), rank(n, 1) { iota(root.begin(), root.end(), 0); }

6

7 int find(int x) {

8 while (x != root[x])

9 x = root[x] = root[root[x]];

10 return x;

11 }

12

13 void join(int x, int y) {

14 x = find(x), y = find(y);

15 if (x != y) {

16 if (rank[x] > rank[y]) swap(x, y);

17 root[x] = y;

18 rank[y] += 1;

19 n--;

20 }

21 }

22

23 int count() { return n; }

24 bool connected(int x, int y) { return find(x) == find(y); }

25 };

26

27 class Solution {

28 typedef vector<int> edge;

29

30 int get_mst(int n, const vector<edge> &edges, int blockedge, int pre_edge = -1) {

31 UnionFind uf(n);

32 int weight = 0;

33 if (pre_edge != -1) {

34 weight += edges[pre_edge][2];

35 uf.join(edges[pre_edge][0], edges[pre_edge][1]);

36 }

37 for (int i = 0; i < edges.size() && uf.count() != 1; ++i) {

38 }

39 if (i == blockedge) continue;

40 const auto &e = edges[i];

41 if (uf.connected(e[0], e[1])) continue;

42 uf.join(e[0], e[1]);

43 weight += e[2];

44 }

45

46 return uf.count() == 1 ? weight : 1e9 + 7;

47 } public : vector<vector<int>> findCriticalAndPseudoCriticalEdges(int n, vector<edge> &edges) {

48 for (int i = 0; i < edges.size(); i++)

49 edges[i].push_back(i);

50 sort(edges.begin(), edges.end(), [](const edge &a, const edge &b) { return a[2] < b[2]; });

51 int origin_mst = get_mst(n, edges, -1);

52

53 vector<int> critical, non_critical;

54 for (int i = 0; i < edges.size(); i++) {

55 if (origin_mst < get_mst(n, edges, i))

56 critical.push_back(edges[i][3]);

57 else if (origin_mst == get_mst(n, edges, -1, i))

58 non_critical.push_back(edges[i][3]);

59 }

60 return {critical, non_critical};

61 }

62 }

63 ;