952. Largest Component Size by Common Factor - LeetCode Solution

class Solution:
    def largestComponentSize(self, A: List[int]) -> int:
        pri = []
        for x in range(2, int(max(A)**0.5)+1):
            for y in pri:
                if x % y == 0:
                    break
            else:
                pri.append(x)
    
        factors = collections.defaultdict(list)         # compute factors of each 'a'
        for a in A:
            x = a
            for p in pri:
                if p*p > x:
                    break
                if x % p == 0:
                    factors[a].append(p)
                    while x % p == 0:
                        x //= p
            if x > 1:                                   # a new prime found
                factors[a].append(x)
                pri.append(x)
                
        pri = list(set(pri))
        n = len(pri)
        p2i = {p: i for i,p in enumerate(pri)}       # prime to index
        
        parent = [i for i in range(n)]                  # union-find on primes
        
        def find(i):
            if i != parent[i]:
                parent[i] = find(parent[i])
            return parent[i]
        
        def union(i,j):
            pi, pj = find(i), find(j)
            if pi != pj:
                parent[pi] = pj

        for a in A:
            if factors[a]:
                p0 = factors[a][0]
                for p in factors[a][1:]:                # link two primes if they are factors of 'a'
                    union(p2i[p0], p2i[p])
        
        count = collections.Counter(find(p2i[factors[a][0]]) for a in A if factors[a])      # each 'a' corresponds to a prime index
        return max(count.values())