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

24 class Solution {

25 public:

26 int numberOfGoodPaths(vector<int> &vals, vector<vector<int>> &edges) {

27 int n = vals.size();

28 map<int, vector<int>> valuesToNodes;

29 vector<vector<int>> adj(n);

30 UnionFind uf(n);

31

32 for (auto &edge : edges) {

33 adj[edge[0]].push_back(edge[1]);

34 adj[edge[1]].push_back(edge[0]);

35 }

36

37 for (int i = 0; i < n; i++)

38 valuesToNodes[vals[i]].push_back(i);

39

40 int goodPaths = 0;

41 for (auto &[value, nodes] : valuesToNodes) {

42 for (int node : nodes) {

43 for (int neighbor : adj[node]) {

44 if (vals[node] >= vals[neighbor]) {

45 uf.join(node, neighbor);

46 }

47 }

48 }

49 unordered_map<int, int> group;

50 for (int u : nodes)

51 group[uf.find(u)]++;

52 for (auto &[_, size] : group)

53 goodPaths += (size * (size + 1) / 2);

54 }

55 return goodPaths;

56 }

57 };