Checking Existence of Edge Length Limited Paths LeetCode solution in C++

1
class Solution {
2
 public:
3
  vector<bool> distanceLimitedPathsExist(int length, vector<vector<int>> &adjList,
4
                                         vector<vector<int>> &queries) {
5
    vector<int> parent(length);
6
    vector<int> rank(length);
7
    vector<int> weight(length);
8
    for (int i = 0; i < length; i++) parent[i] = i;
9

10
    sort(adjList.begin(), adjList.end(),
11
         [](const vector<int> &a, const vector<int> &b) { return a[2] < b[2]; });
12
    for (vector<int> &edge : adjList) unionByRank(edge[0], edge[1], edge[2], parent, rank, weight);
13

14
    vector<bool> answer;
15
    for (vector<int> &query : queries)
16
      answer.push_back(isConnectedAndWithinLimit(query[0], query[1], query[2], parent, weight));
17

18
    return answer;
19
  }
20

21
  bool isConnectedAndWithinLimit(int p, int q, int limit, vector<int> &parent,
22
                                 vector<int> &weight) {
23
    return find(p, limit, parent, weight) == find(q, limit, parent, weight);
24
  }
25

26
  int find(int x, int limit, vector<int> &parent, vector<int> &weight) {
27
    while (x != parent[x]) {
28
      if (weight[x] >= limit) {
29
        break;
30
      }
31
      x = parent[x];
32
    }
33
    return x;
34
  }
35

36
  void unionByRank(int x, int y, int limit, vector<int> &parent, vector<int> &rank,
37
                   vector<int> &weight) {
38
    int x_ref = find(x, INT_MAX, parent, weight);
39
    int y_ref = find(y, INT_MAX, parent, weight);
40
    if (x_ref != y_ref) {
41
      if (rank[x_ref] < rank[y_ref]) {
42
        parent[x_ref] = y_ref;
43
        weight[x_ref] = limit;
44
      } else {
45
        parent[y_ref] = x_ref;
46
        weight[y_ref] = limit;
47
        if (rank[x_ref] == rank[y_ref]) {
48
          rank[x_ref]++;
49
        }
50
      }
51
    }
52
  }
53
};