Find Critical and Pseudo-Critical Edges in Minimum Spanning Tree LeetCode solution in C++

1
class UnionFind {
2
 private:
3
  vector<int> parent_;
4
  vector<int> rank_;
5
  int sets_;
6

7
 public:
8
  UnionFind(int n) {
9
    init(n);
10
  }
11
  void init(int n) {
12
    sets_ = n;
13
    parent_.resize(n);
14
    rank_.resize(n);
15
    iota(parent_.begin(), parent_.end(), 0);
16
    fill(rank_.begin(), rank_.end(), 1);
17
  }
18
  int find(int u) {
19
    return parent_[u] == u ? u : parent_[u] = find(parent_[u]);
20
  }
21
  bool join(int u, int v) {
22
    u = find(u);
23
    v = find(v);
24

25
    if (u == v) {
26
      return false;
27
    }
28

29
    if (rank_[u] < rank_[v]) {
30
      swap(u, v);
31
    }
32

33
    rank_[u] += rank_[v];
34
    parent_[v] = u;
35

36
    sets_--;
37
    return true;
38
  }
39
  bool united() {
40
    return sets_ == 1;
41
  }
42
};
43
class Solution {
44
 public:
45
  vector<int> edges_idx_;
46
  int kruskal(const int n, const int removed_edge_idx, const int init_edge_idx,
47
              const vector<vector<int>> &edges) {
48
    UnionFind graph(n);
49
    int edges_size = edges.size();
50
    int total = 0;
51

52
    if (init_edge_idx != -1) {
53
      graph.join(edges[init_edge_idx][0], edges[init_edge_idx][1]);
54
      total += edges[init_edge_idx][2];
55
    }
56

57
    for (int i = 0; i < edges_size; ++i) {
58
      int edge_idx = edges_idx_[i];
59

60
      if (edge_idx == removed_edge_idx) {
61
        continue;
62
      }
63

64
      int u = edges[edge_idx][0];
65
      int v = edges[edge_idx][1];
66

67
      if (graph.join(u, v)) {
68
        total += edges[edge_idx][2];
69
      }
70
    }
71

72
    return graph.united() ? total : INT_MAX;
73
  }
74
  vector<vector<int>> findCriticalAndPseudoCriticalEdges(int n, vector<vector<int>> &edges) {
75
    int edges_size = edges.size();
76

77
    edges_idx_.resize(edges_size);
78

79
    iota(edges_idx_.begin(), edges_idx_.end(), 0);
80

81
    sort(edges_idx_.begin(), edges_idx_.end(),
82
         [&edges](int a, int b) { return edges[a][2] < edges[b][2]; });
83

84
    int mst = kruskal(n, -1, -1, edges);
85

86
    vector<int> critical;
87
    vector<int> psudo;
88

89
    for (int i = 0; i < edges_size; ++i) {
90
      int edge_idx = edges_idx_[i];
91
      int total = kruskal(n, -1, edge_idx, edges);
92

93
      if (total == mst) {
94
        total = kruskal(n, edge_idx, -1, edges);
95
        if (total > mst) {
96
          critical.push_back(edge_idx);
97
        } else {
98
          psudo.push_back(edge_idx);
99
        }
100
      }
101
    }
102

103
    return {critical, psudo};
104
  }
105
};