leetcode

Solution to some Leetcode problems written in C++

git clone git://git.dimitrijedobrota.com/leetcode.git

2421.cpp (1550B)

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 };