Largest Component Size by Common Factor LeetCode solution in Java

1
class Solution {
2
  public int largestComponentSize(int[] nums) {
3
    int maxNum = getMaxNum(nums);
4
    Map<Integer, Integer> numToFirstPrimeFactor = new HashMap<>();
5
    DisjointSet ds = new DisjointSet(maxNum + 1);
6
    for (int num : nums) {
7
      if (num == 1) {
8
        continue;
9
      }
10

11
      List<Integer> primeFactors = getPrimeFactors(num);
12
      int firstPrimeFactor = primeFactors.get(0);
13
      numToFirstPrimeFactor.put(num, firstPrimeFactor);
14

15
      for (int i = 1; i < primeFactors.size(); i++) {
16
        ds.union(primeFactors.get(i - 1), primeFactors.get(i));
17
      }
18
    }
19

20
    Map<Integer, Integer> componentToSize = new HashMap<>();
21
    int maxSize = 0;
22
    for (int num : nums) {
23
      if (num == 1) {
24
        continue;
25
      }
26

27
      int firstPrimeFactor = numToFirstPrimeFactor.get(num);
28
      int component = ds.find(firstPrimeFactor);
29
      int size = componentToSize.getOrDefault(component, 0);
30
      componentToSize.put(component, ++size);
31
      maxSize = Math.max(maxSize, size);
32
    }
33
    return maxSize;
34
  }
35

36
  public int getMaxNum(int[] nums) {
37
    int maxNum = 0;
38
    for (int num : nums) {
39
      maxNum = Math.max(maxNum, num);
40
    }
41
    return maxNum;
42
  }
43

44
  public List<Integer> getPrimeFactors(int num) {
45
    List<Integer> primeFactors = new ArrayList<>();
46

47
    // even prime factor i.e. 2
48
    if ((num & 1) == 0) {
49
      primeFactors.add(2);
50

51
      do {
52
        num >>= 1;
53
      } while ((num & 1) == 0);
54
    }
55

56
    // odd prime factors
57
    int primeFactor = 3;
58
    while (num != 1 && primeFactor * primeFactor <= num) {
59
      if (num % primeFactor == 0) {
60
        primeFactors.add(primeFactor);
61

62
        do {
63
          num /= primeFactor;
64
        } while (num % primeFactor == 0);
65
      }
66
      primeFactor += 2;
67
    }
68

69
    // num is prime
70
    if (num != 1) {
71
      primeFactors.add(num);
72
    }
73
    return primeFactors;
74
  }
75
}
76

77
class DisjointSet {
78
  int[] root;
79
  int[] rank;
80

81
  public DisjointSet(int size) {
82
    root = new int[size];
83
    rank = new int[size];
84
    for (int i = 0; i < size; i++) {
85
      root[i] = i;
86
      rank[i] = 1;
87
    }
88
  }
89

90
  public int find(int x) {
91
    while (x != root[x]) {
92
      root[x] = root[root[x]];
93
      x = root[x];
94
    }
95
    return x;
96
  }
97

98
  public void union(int x, int y) {
99
    int rootX = find(x);
100
    int rootY = find(y);
101

102
    if (rootX == rootY) {
103
      return;
104
    }
105

106
    if (rank[rootX] > rank[rootY]) {
107
      root[rootY] = rootX;
108
    } else if (rank[rootX] < rank[rootY]) {
109
      root[rootX] = rootY;
110
    } else {
111
      root[rootY] = rootX;
112
      rank[rootX]++;
113
    }
114
  }
115
}