Largest Component Size by Common Factor LeetCode solution in Python

1
class UnionFind:
2

3
    def __init__(self, n):
4
        self.parent = list(range(n))
5
        self.rank = [1] * n
6

7
    def find(self, p):
8
        if p != self.parent[p]:
9
            self.parent[p] = self.find(self.parent[p])
10
        return self.parent[p]
11

12
    def union(self, p, q):
13
        prt, qrt = self.find(p), self.find(q)
14
        if prt == qrt:
15
            return False
16
        if self.rank[prt] > self.rank[qrt]:
17
            prt, qrt = qrt, prt
18
        self.parent[prt] = qrt
19
        self.rank[qrt] += self.rank[prt]
20
        return True
21

22

23
class Solution:
24
    def largestComponentSize(self, A: List[int]) -> int:
25
        m = max(A)
26
        uf = UnionFind(m + 1)
27
        seen = set(A)
28

29
        # modified sieve of eratosthenes
30
        sieve = [1] * (m + 1)
31
        sieve[0] = sieve[1] = 0
32
        for k in range(m // 2 + 1):
33
            if sieve[k]:
34
                prev = k if k in seen else 0
35
                for x in range(2 * k, m + 1, k):
36
                    sieve[x] = 0
37
                    if x in seen:
38
                        if prev:
39
                            uf.union(prev, x)
40
                        else:
41
                            prev = x
42
        return max(uf.rank)