LC553 Largest Component Size by Common Factor

class Solution:
    # time O(NlgM) where M is factors and numers .  space O(M+N) N for map, M for UF
        uf = UF(max_val)
        #step1, attribute num in A to all groups that lead by factors
        for num in A:
            factor_max = int(sqrt(num))
            for fac in range(2,factor_max+1):
                if num % fac == 0:
                    uf.union(num,fac)
                    uf.union(num,num//fac)
        #step2 count size of group one by one and get the max cnt 
        size_max = 0 
        group_cnt = collections.defaultdict(int)
        for num in A:
            group_id = uf.find(num)
            group_cnt[group_id] += 1
            size_max = max(size_max,group_cnt[group_id])
        return size_max
        
class UF:
    def __init__(self,size):
        # each node is an independent component
        self.parent = [i for i in range(size+1)]
        #keep size of each component 
        self.size = [1] * (size+1)
    
    def find(self, x):
        """return component id that element x belong to """
        while x != self.parent[x]:
            self.parent[x] = self.parent[self.parent[x]] # optimization using path compression
            # if x != self.parent[x]:
            #   self.parent[x] = self.find(self.parent[x]) 
            # return self.parent[x]
            x = self.parent[x]
        return x 
    
    def union(self, x, y):
        """merge two components that x,y belongs to sepereately, and return merged component """
        r_x = self.find(x)
        r_y = self.find(y)
        if r_x != r_y:
            self.parent[r_x] = r_y
            self.size[r_y] += self.size[r_x]
        return r_y