Count the Number of Ideal Arrays LeetCode solution in Java

1
class Solution {
2
  int M = (int) 1e9 + 7;
3

4
  public int idealArrays(int n, int maxValue) {
5
    long ans = 0;
6
    int N = n + maxValue;
7
    long[] inv = new long[N];
8
    long[] fact = new long[N];
9
    long[] factinv = new long[N];
10
    inv[1] = fact[0] = fact[1] = factinv[0] = factinv[1] = 1;
11
    for (int i = 2; i < N; i++) { // mod inverse
12
      inv[i] = M - M / i * inv[M % i] % M;
13
      fact[i] = fact[i - 1] * i % M;
14
      factinv[i] = factinv[i - 1] * inv[i] % M;
15
    }
16
    for (int i = 1; i <= maxValue; i++) {
17
      int tmp = i;
18
      Map<Integer, Integer> map = new HashMap<>();
19
      for (int j = 2; j * j <= tmp; j++) {
20
        while (tmp % j == 0) { // prime factorization.
21
          tmp /= j;
22
          map.merge(j, 1, Integer::sum);
23
        }
24
      }
25
      if (tmp > 1) {
26
        map.merge(tmp, 1, Integer::sum);
27
      }
28
      long gain = 1;
29
      for (int val : map.values()) { // arranges all the primes.
30
        gain *= comb(n + val - 1, val, fact, factinv);
31
        gain %= M;
32
      }
33
      ans += gain;
34
      ans %= M;
35
    }
36

37
    return (int) ans;
38
  }
39

40
  private long comb(int a, int b, long[] fact, long[] factinv) {
41
    return fact[a] * factinv[b] % M * factinv[a - b] % M;
42
  }
43
}