leetcode

0 class UnionFind {

1 int n, cnt = n;

2 vector<int> root, size;

3

4 public:

5 UnionFind(int n) : n(n), root(n), size(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 (size[x] > size[y]) swap(x, y);

17 root[x] = y;

18 size[y] += size[x];

19 cnt--;

20 }

21 }

22

23 int count() { return cnt; }

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

25 };

26

27 class Solution {

28 public:

29 vector<bool> distanceLimitedPathsExist(int n, vector<vector<int>> &edges, vector<vector<int>> &queries) {

30 vector<bool> res(queries.size());

31 UnionFind uf(n);

32

33 for (int i = 0; i < queries.size(); ++i)

34 queries[i].push_back(i);

35

36 const auto cmp = [](auto &a, auto &b) { return a[2] < b[2]; };

37 sort(begin(queries), end(queries), cmp);

38 sort(begin(edges), end(edges), cmp);

39

40 int i = 0;

41 for (auto &q : queries) {

42 for (; i < edges.size() && edges[i][2] < q[2]; i++)

43 uf.join(edges[i][0], edges[i][1]);

44 res[q[3]] = uf.connected(q[0], q[1]);

45 }

46 return res;

47 }

48 };