Largest Component Size By Common Factor

import java.util.*;
import java.util.stream.*;
import java.math.*;

public class Solution {

    //并查集
    int MAX=100000;

    //用不加权并查集的话，把count去掉就行，代码简单一点
    int[] count=new int[MAX],find = new int[MAX];

    //这个是加权并查集的方法
    public int count(int a){
        return count[find(a)];
    }

    public int find(int a){
        while(find[a]!=a)a=find[a];
        return a;
    }

    public void union(int a,int b){
        a=find(a);
        b=find(b);

        if(a==b)return;

        if(count[a]>count[b]){
            count[a]+=count[b];
            find[b] = a;
        }else{
            count[b]+=count[a];
            find[a] = b;
        }
    }

    public int largestComponentSize(int[] A) {
        //初始化并查集
        for (int i = 0; i < MAX; i++) {
            find[i] = i;
            count[i] = 1;
        }

        Map<Integer,Integer> factorMap = new HashMap<>();

        for (int i = 0; i < A.length; i++) {
            double sqrt=Math.sqrt(A[i]);
            for (int j = 2; j <=sqrt; j++) {
                if(A[i] % j==0){
                    unionMap(factorMap,j,i);
                    unionMap(factorMap,A[i]/j,i);
                }
            }

            unionMap(factorMap,A[i],i);
        }

        int max = 1;

        for (int i = 0; i < A.length; i++) {
            max = Math.max(max,count(i));
        }

        return max;
    }

    public void unionMap(Map<Integer,Integer> factorMap,int factor,int cur){
        if(factorMap.containsKey(factor)){
            union(factorMap.get(factor),cur);
        }else{
            factorMap.put(factor,cur);
        }
    }
}