Graph Connectivity With Threshold LeetCode solution in Python

1
# Runtime: 881 ms (Top 97.89%) | Memory: 49.3 MB (Top 35.92%)
2
class Solution:
3
    def areConnected(
4
        self, n: int, threshold: int, queries: List[List[int]]
5
    ) -> List[bool]:
6
        parent = list(range(n + 1))
7

8
        def find(i):
9
            if parent[i] != i:
10
                parent[i] = find(parent[i])
11
            return parent[i]
12

13
        def union(i, j):
14
            parent[find(i)] = find(j)
15

16
        if not threshold:
17
            return [True] * len(queries)
18
        for i in range(1, n + 1):
19
            for j in range(2 * i, n + 1, i):
20
                if i > threshold:
21
                    union(i, j)
22
        return [find(i) == find(j) for i, j in queries]