largest_components.py

class Solution:
    def largestComponentSize(self, A):
        return self.figure_components(A, self.sieve())

    def sieve(self):
        range_max = 100001
        sieve_array = [-1] * range_max
        for i in range(2, range_max):
            for j in range(i, range_max, i):
                if sieve_array[j] == -1: sieve_array[j] = i
        return sieve_array
    
    def figure_components(self, arr, sieve_array):
        prime_to_idx = {}
        tree = [-1] * len(arr)
        for i, n in enumerate(arr):
            tmp_n = n
            while tmp_n > 1:
                highest_prime_fac = sieve_array[tmp_n]
                tmp_n //= highest_prime_fac
                if highest_prime_fac in prime_to_idx:
                    self.onion(tree, prime_to_idx[highest_prime_fac], i)
                else:
                    prime_to_idx[highest_prime_fac] = i
        return -1 * min(tree)
                    
    def onion(self, t, a, b):
        a, b = self.root(t, a), self.root(t, b)
        if a == b: return
        if t[a] > t[b]: a, b = b, a
        t[a] += t[b]
        t[b] = a
    
    def root(self, t, a):
        if t[a] < 0: return a
        t[a] = self.root(t, t[a])
        return t[a]