Skip to content

1735. Count Ways to Make Array With Product 👍

  • Time: $O(10^4 + q\sqrt{n})$
  • Space: $O(10^4 + q)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
class Solution {
 public:
  vector<int> waysToFillArray(vector<vector<int>>& queries) {
    constexpr int kMax = 10000;
    constexpr int kMaxFreq = 13;  // 2^13 = 8192 < kMax
    const vector<int> minPrimeFactors = sieveEratosthenes(kMax + 1);
    const auto [fact, invFact] = getFactAndInvFact(kMax + kMaxFreq - 1);
    vector<int> ans;

    for (const vector<int>& query : queries) {
      const int n = query[0];
      const int k = query[1];
      int res = 1;
      for (const auto& [_, freq] : getPrimeFactorsCount(k, minPrimeFactors))
        res = static_cast<long>(res) * nCk(n - 1 + freq, freq, fact, invFact) %
              kMod;
      ans.push_back(res);
    }

    return ans;
  }

 private:
  static constexpr int kMod = 1'000'000'007;

  // Gets the minimum prime factor of i, where 1 < i <= n.
  vector<int> sieveEratosthenes(int n) {
    vector<int> minPrimeFactors(n + 1);
    iota(minPrimeFactors.begin() + 2, minPrimeFactors.end(), 2);
    for (int i = 2; i * i < n; ++i)
      if (minPrimeFactors[i] == i)  // `i` is prime.
        for (int j = i * i; j < n; j += i)
          minPrimeFactors[j] = min(minPrimeFactors[j], i);
    return minPrimeFactors;
  }

  unordered_map<int, int> getPrimeFactorsCount(
      int num, const vector<int>& minPrimeFactors) {
    unordered_map<int, int> count;
    while (num > 1) {
      const int divisor = minPrimeFactors[num];
      while (num % divisor == 0) {
        num /= divisor;
        ++count[divisor];
      }
    }
    return count;
  }

  pair<vector<long>, vector<long>> getFactAndInvFact(int n) {
    vector<long> fact(n + 1);
    vector<long> invFact(n + 1);
    vector<long> inv(n + 1);
    fact[0] = invFact[0] = 1;
    inv[0] = inv[1] = 1;
    for (int i = 1; i <= n; ++i) {
      if (i >= 2)
        inv[i] = kMod - kMod / i * inv[kMod % i] % kMod;
      fact[i] = fact[i - 1] * i % kMod;
      invFact[i] = invFact[i - 1] * inv[i] % kMod;
    }
    return {fact, invFact};
  }

  int nCk(int n, int k, const vector<long>& fact, const vector<long>& invFact) {
    return fact[n] * invFact[k] % kMod * invFact[n - k] % kMod;
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
class Solution {
  public int[] waysToFillArray(int[][] queries) {
    final int kMax = 10_000;
    final int kMaxFreq = 13; // 2^13 = 8192 < kMax
    final int[] minPrimeFactors = sieveEratosthenes(kMax + 1);
    final long[][] factAndInvFact = getFactAndInvFact(kMax + kMaxFreq - 1);
    final long[] fact = factAndInvFact[0];
    final long[] invFact = factAndInvFact[1];
    int[] ans = new int[queries.length];

    for (int i = 0; i < queries.length; ++i) {
      final int n = queries[i][0];
      final int k = queries[i][1];
      int res = 1;
      for (final int freq : getPrimeFactorsCount(k, minPrimeFactors).values())
        res = (int) ((long) res * nCk(n - 1 + freq, freq, fact, invFact) % kMod);
      ans[i] = res;
    }

    return ans;
  }

  private static final int kMod = 1_000_000_007;

  // Gets the minimum prime factor of i, where 1 < i <= n.
  private int[] sieveEratosthenes(int n) {
    int[] minPrimeFactors = new int[n + 1];
    for (int i = 2; i <= n; ++i)
      minPrimeFactors[i] = i;
    for (int i = 2; i * i < n; ++i)
      if (minPrimeFactors[i] == i) // `i` is prime.
        for (int j = i * i; j < n; j += i)
          minPrimeFactors[j] = Math.min(minPrimeFactors[j], i);
    return minPrimeFactors;
  }

  private Map<Integer, Integer> getPrimeFactorsCount(int num, int[] minPrimeFactors) {
    Map<Integer, Integer> count = new HashMap<>();
    while (num > 1) {
      final int divisor = minPrimeFactors[num];
      while (num % divisor == 0) {
        num /= divisor;
        count.put(divisor, count.merge(divisor, 1, Integer::sum));
      }
    }
    return count;
  }

  private long[][] getFactAndInvFact(int n) {
    long[] fact = new long[n + 1];
    long[] invFact = new long[n + 1];
    long[] inv = new long[n + 1];
    fact[0] = invFact[0] = 1;
    inv[0] = inv[1] = 1;
    for (int i = 1; i <= n; ++i) {
      if (i >= 2)
        inv[i] = kMod - kMod / i * inv[kMod % i] % kMod;
      fact[i] = fact[i - 1] * i % kMod;
      invFact[i] = invFact[i - 1] * inv[i] % kMod;
    }
    return new long[][] {fact, invFact};
  }

  private int nCk(int n, int k, long[] fact, long[] invFact) {
    return (int) (fact[n] * invFact[k] % kMod * invFact[n - k] % kMod);
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
class Solution:
  def waysToFillArray(self, queries: list[list[int]]) -> list[int]:
    kMod = 1_000_000_007
    kMax = 10_000
    minPrimeFactors = self._sieveEratosthenes(kMax + 1)

    @functools.lru_cache(None)
    def fact(i: int) -> int:
      return 1 if i <= 1 else i * fact(i - 1) % kMod

    @functools.lru_cache(None)
    def inv(i: int) -> int:
      return pow(i, kMod - 2, kMod)

    @functools.lru_cache(None)
    def nCk(n: int, k: int) -> int:
      return fact(n) * inv(fact(k)) * inv(fact(n - k)) % kMod

    ans = []

    for n, k in queries:
      res = 1
      for freq in self._getPrimeFactorsCount(k, minPrimeFactors).values():
        res = res * nCk(n - 1 + freq, freq) % kMod
      ans.append(res)

    return ans

  def _sieveEratosthenes(self, n: int) -> list[int]:
    """Gets the minimum prime factor of i, where 1 < i <= n."""
    minPrimeFactors = [i for i in range(n + 1)]
    for i in range(2, int(n**0.5) + 1):
      if minPrimeFactors[i] == i:  # `i` is prime.
        for j in range(i * i, n, i):
          minPrimeFactors[j] = min(minPrimeFactors[j], i)
    return minPrimeFactors

  def _getPrimeFactorsCount(self, num: int, minPrimeFactors: list[int]) -> dict[int, int]:
    count = collections.Counter()
    while num > 1:
      divisor = minPrimeFactors[num]
      while num % divisor == 0:
        num //= divisor
        count[divisor] += 1
    return count