3532. Path Existence Queries in a Graph I | 1659
initial union find version:
1class DSU:
2 def __init__(self, n):
3 self.rank = [0] * n
4 self.parent = [i for i in range(n)]
5
6 def find(self, i):
7 if self.parent[i] != i:
8 self.parent[i] = self.find(self.parent[i])
9
10 return self.parent[i]
11
12 def union(self, i, j):
13 i_parent = self.find(i)
14 j_parent = self.find(j)
15
16 if i_parent == j_parent:
17 return
18
19 i_parent_rank = self.rank[i_parent]
20 j_parent_rank = self.rank[j_parent]
21
22 if i_parent_rank < j_parent_rank:
23 self.parent[i_parent] = j_parent
24 elif i_parent_rank > j_parent_rank:
25 self.parent[j_parent] = i_parent
26 else:
27 self.parent[j_parent] = i_parent
28 self.rank[i_parent] += 1
29
30
31class Solution:
32 def pathExistenceQueries(
33 self, n: int, nums: List[int], maxDiff: int, queries: List[List[int]]
34 ) -> List[bool]:
35 # feels like a perfect problem for union find
36
37 uf = DSU(n)
38
39 l, r = 0, 0
40 while r < n:
41 while r < n and abs(nums[l] - nums[r]) <= maxDiff:
42 uf.union(l, r)
43 r += 1
44 while r < n and abs(nums[l] - nums[r]) > maxDiff:
45 l += 1
46
47 return [uf.find(a) == uf.find(b) for a, b in queries]
same asymptotic but better code:
1class DSU:
2 def __init__(self, n):
3 self.rank = [0] * n
4 self.parent = [i for i in range(n)]
5
6 def find(self, i):
7 if self.parent[i] != i:
8 self.parent[i] = self.find(self.parent[i])
9
10 return self.parent[i]
11
12 def union(self, i, j):
13 i_parent = self.find(i)
14 j_parent = self.find(j)
15
16 if i_parent == j_parent:
17 return
18
19 i_parent_rank = self.rank[i_parent]
20 j_parent_rank = self.rank[j_parent]
21
22 if i_parent_rank < j_parent_rank:
23 self.parent[i_parent] = j_parent
24 elif i_parent_rank > j_parent_rank:
25 self.parent[j_parent] = i_parent
26 else:
27 self.parent[j_parent] = i_parent
28 self.rank[i_parent] += 1
29
30
31class Solution:
32 def pathExistenceQueries(
33 self, n: int, nums: List[int], maxDiff: int, queries: List[List[int]]
34 ) -> List[bool]:
35 uf = DSU(n)
36
37 # Only need to connect adjacent indices when their values are close enough.
38 for i in range(n - 1):
39 if nums[i + 1] - nums[i] <= maxDiff:
40 uf.union(i, i + 1)
41
42 return [uf.find(u) == uf.find(v) for u, v in queries]
DSU is actually not required for this problem:
1class Solution:
2 def pathExistenceQueries(
3 self, n: int, nums: List[int], maxDiff: int, queries: List[List[int]]
4 ) -> List[bool]:
5 comp = [0] * n # component id
6 comp[0] = 0
7 for i in range(1, n):
8 if nums[i] - nums[i-1] <= maxDiff:
9 comp[i] = comp[i-1]
10 else:
11 comp[i] = i # new component
12
13 return [comp[u] == comp[v] for u, v in queries]