- Give each element a component representative.
- Merge representatives when a connection is accepted.
- Answer connectivity or component queries from the compressed representatives.
Code notes
- 47 lines of Python from the credited upstream file number-of-good-paths.py.
- The implementation visibly relies on sequence storage, hash lookup.
- No explicit loop blocks detected.
Complexity
Account for every find and union operation; with path compression and ranked merging, the amortized cost is nearly constant per operation.
Check the problem constraints before deciding whether this complexity will pass.
Use this to learn the idea, then write your own version.
123 4import collections5 6 7class UnionFind(object): 8 def __init__(self, vals):9 self.set = range(len(vals))10 self.rank = [0]*len(vals)11 self.cnt = [collections.Counter({v:1}) for v in vals] 12 13 def find_set(self, x):14 stk = []15 while self.set[x] != x: 16 stk.append(x)17 x = self.set[x]18 while stk:19 self.set[stk.pop()] = x20 return x21 22 def union_set(self, x, y, v): 23 x, y = self.find_set(x), self.find_set(y)24 if x == y:25 return 0 26 if self.rank[x] > self.rank[y]: 27 x, y = y, x28 self.set[x] = self.set[y]29 if self.rank[x] == self.rank[y]:30 self.rank[y] += 131 cx, cy = self.cnt[x][v], self.cnt[y][v] 32 self.cnt[y] = collections.Counter({v:cx+cy}) 33 return cx*cy 34 35 3637class Solution(object):38 def numberOfGoodPaths(self, vals, edges):39 """40 :type vals: List[int]41 :type edges: List[List[int]]42 :rtype: int43 """44 edges.sort(key=lambda x: max(vals[x[0]], vals[x[1]]))45 uf = UnionFind(vals)46 return len(vals)+sum(uf.union_set(i, j, max(vals[i], vals[j])) for i, j in edges)47