Problem statement

https://leetcode.com/problems/count-ways-to-make-array-with-product/

Solution

The idea is to ask question, given $p_i^{\alpha_i}$, what is the number of ways to allocate this prime factor. In fact this is stars and bars problem and it can be done in $C_{n+\alpha_i - 1}^{n-1}$ ways to put them into $n$ groups. So, what we can do is to factorize number and then multiply number of cases.

Complexity

It is O(M*log M) to calculate factorials and arr, then for each query we can answer it in O(log M), so total time complexity is O(M*log M + Q * log M), where Q is number of queries.

Code

class Solution:
    MOD, M = 10**9 + 7, 10**4 + 20
    F = [1]*(M+1)
    for i in range(1, M+1):
        F[i] = (F[i-1] * i) % MOD

    I = [1]*M + [pow(F[M], MOD - 2, MOD)]
    for i in range(M-1, 0, -1):
        I[i] = I[i+1]*(i+1) % MOD

    arr = [0]*(M+1)
    for i in range(2, M+1):
        for j in range(i, M+1, i):
            if arr[j] == 0: arr[j] = i
    
    def waysToFillArray(self, Q):
        ans = []
        for n, k in Q:
            primes = []
            while self.arr[k] != 0:
                primes += [self.arr[k]]
                k //= self.arr[k]
            res = 1
            for x in Counter(primes).values():
                res = (res * self.F[x + n - 1] * self.I[n-1] * self.I[x]) % self.MOD
            ans += [res]
            
        return ans