philosyang.com

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]

#Graph #Union-Find #Python