Graph Connectivity With Threshold LeetCode solution in C++

1
// DSU Class Template
2
class DSU {
3
  vector<int> parent, size;
4

5
 public:
6
  DSU(int n) {
7
    for (int i = 0; i <= n; i++) {
8
      parent.push_back(i);
9
      size.push_back(1);
10
    }
11
  }
12

13
  int findParent(int num) {
14
    if (parent[num] == num) return num;
15
    return parent[num] = findParent(parent[num]);
16
  }
17

18
  void unionBySize(int u, int v) {
19
    int parU = findParent(u);
20
    int parV = findParent(v);
21

22
    if (parU == parV) return;
23

24
    if (size[parU] < size[parV]) {
25
      parent[parU] = parV;
26
      size[parV] += size[parU];
27
    } else {
28
      parent[parV] = parU;
29
      size[parU] += size[parV];
30
    }
31
  }
32
};
33

34
class Solution {
35
 public:
36
  vector<bool> areConnected(int n, int threshold, vector<vector<int>> &queries) {
37
    DSU dsu(n);
38

39
    // Make connections with its multiple
40
    for (int i = threshold + 1; i <= n; i++) {
41
      int j = 1;
42
      while (i * j <= n) {
43
        dsu.unionBySize(i, i * j);
44
        j++;
45
      }
46
    }
47

48
    vector<bool> res;
49
    for (auto &query : queries) {
50
      int u = query[0];
51
      int v = query[1];
52

53
      // Check if both cities belong to same component or not
54
      if (dsu.findParent(u) == dsu.findParent(v))
55
        res.push_back(true);
56
      else
57
        res.push_back(false);
58
    }
59

60
    return res;
61
  }
62
};