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

1
class Solution {
2
  static class UnionFind {
3
    int[] parent;
4
    int[] rank;
5
    int comp = 0;
6

7
    UnionFind(int n) {
8
      parent = new int[n];
9
      rank = new int[n];
10
      comp = n;
11
      for (int i = 0; i < n; i++) {
12
        parent[i] = i;
13
        rank[i] = 0;
14
      }
15
    }
16

17
    int find(int x) {
18
      if (parent[x] == x) {
19
        return x;
20
      } else {
21
        parent[x] = find(parent[x]);
22
        return parent[x];
23
      }
24
    }
25

26
    boolean union(int X, int Y) {
27
      int x = find(X);
28
      int y = find(Y);
29
      if (x == y) {
30
        return false;
31
      }
32

33
      if (rank[x] < rank[y]) {
34
        parent[x] = y;
35
      } else if (rank[y] < rank[x]) {
36
        parent[y] = x;
37
      } else {
38
        parent[y] = x;
39
        rank[x]++;
40
      }
41
      comp--;
42
      return true;
43
    }
44

45
    boolean isConnected() {
46
      return comp == 1;
47
    }
48
  }
49

50
  static class Edge implements Comparable<Edge> {
51
    int u;
52
    int v;
53
    int wt;
54

55
    Edge(int u, int v, int wt) {
56
      this.u = u;
57
      this.v = v;
58
      this.wt = wt;
59
    }
60

61
    public int compareTo(Edge o) {
62
      return this.wt - o.wt;
63
    }
64
  }
65

66
  public int buildMST(int n, int[][] edges, int[] edgeSkip, int[] edgePick) {
67
    PriorityQueue<Edge> pq = new PriorityQueue<>();
68

69
    for (int[] edge : edges) {
70
      if (edge == edgeSkip) {
71
        continue;
72
      } else if (edge == edgePick) {
73
        continue;
74
      }
75
      int u = edge[0];
76
      int v = edge[1];
77
      int wt = edge[2];
78
      pq.add(new Edge(u, v, wt));
79
    }
80

81
    UnionFind uf = new UnionFind(n);
82
    int cost = 0;
83

84
    if (edgePick != null) {
85
      uf.union(edgePick[0], edgePick[1]);
86
      cost += edgePick[2];
87
    }
88
    while (pq.size() > 0) {
89
      Edge rem = pq.remove();
90
      if (uf.union(rem.u, rem.v) == true) {
91
        cost += rem.wt;
92
      }
93
    }
94

95
    if (uf.isConnected() == true) {
96
      return cost;
97
    } else {
98
      return Integer.MAX_VALUE;
99
    }
100
  }
101

102
  public List<List<Integer>> findCriticalAndPseudoCriticalEdges(int n, int[][] edges) {
103
    int mstCost = buildMST(n, edges, null, null);
104

105
    ArrayList<Integer> critical = new ArrayList<>();
106
    ArrayList<Integer> pcritical = new ArrayList<>();
107

108
    for (int i = 0; i < edges.length; i++) {
109
      int[] edge = edges[i];
110

111
      int mstCostWithoutEdge = buildMST(n, edges, edge, null);
112
      if (mstCostWithoutEdge > mstCost) {
113
        critical.add(i); // Critical edge index
114
      } else {
115
        int mstCostWithEdge = buildMST(n, edges, null, edge);
116
        if (mstCostWithEdge > mstCost) {
117
          // redundant
118
        } else {
119
          pcritical.add(i); // pseduo critical edge index
120
        }
121
      }
122
    }
123

124
    List<List<Integer>> res = new ArrayList<>();
125
    res.add(critical);
126
    res.add(pcritical);
127
    return res;
128
  }
129
}