2338. Count the Number of Ideal Arrays

Problem Link

Description

---

You are given two integers n and maxValue, which are used to describe an ideal array.

A 0-indexed integer array arr of length n is considered ideal if the following conditions hold:

• Every arr[i] is a value from 1 to maxValue, for 0 <= i < n.
• Every arr[i] is divisible by arr[i - 1], for 0 < i < n.

Return the number of distinct ideal arrays of length n. Since the answer may be very large, return it modulo 109 + 7.

 

Example 1:

Input: n = 2, maxValue = 5
Output: 10
Explanation: The following are the possible ideal arrays:
- Arrays starting with the value 1 (5 arrays): [1,1], [1,2], [1,3], [1,4], [1,5]
- Arrays starting with the value 2 (2 arrays): [2,2], [2,4]
- Arrays starting with the value 3 (1 array): [3,3]
- Arrays starting with the value 4 (1 array): [4,4]
- Arrays starting with the value 5 (1 array): [5,5]
There are a total of 5 + 2 + 1 + 1 + 1 = 10 distinct ideal arrays.

Example 2:

Input: n = 5, maxValue = 3
Output: 11
Explanation: The following are the possible ideal arrays:
- Arrays starting with the value 1 (9 arrays): 
   - With no other distinct values (1 array): [1,1,1,1,1] 
   - With 2nd distinct value 2 (4 arrays): [1,1,1,1,2], [1,1,1,2,2], [1,1,2,2,2], [1,2,2,2,2]
   - With 2nd distinct value 3 (4 arrays): [1,1,1,1,3], [1,1,1,3,3], [1,1,3,3,3], [1,3,3,3,3]
- Arrays starting with the value 2 (1 array): [2,2,2,2,2]
- Arrays starting with the value 3 (1 array): [3,3,3,3,3]
There are a total of 9 + 1 + 1 = 11 distinct ideal arrays.

 

Constraints:

• 2 <= n <= 104
• 1 <= maxValue <= 104

Solution

---

Python3

from math import sqrt
 
class Solution:
    def primesUpTo(self, n):
        primes = set(range(2, n + 1))
        for i in range(2, n):
            if i in primes:
                it = i * 2
                while it <= n:
                    if it in primes:
                        primes.remove(it)
                    it += i
 
        return primes
 
    def getPrimeFactors(self, n, primes):
        ret = {}
        sq = int(math.sqrt(n))
 
        for p in primes:
            if n in primes:
                ret[n] = 1
                break
 
            while n % p == 0:
                ret[p] = ret.get(p, 0) + 1
                n //= p
 
            if n <= 1:
                break
 
        return ret
        
    def idealArrays(self, n: int, maxValue: int) -> int:
        mod = 10**9 + 7
        ret = 0
        primes = self.primesUpTo(maxValue)
        
        for num in range(1, maxValue + 1):
            # find number of arrays that can end with num
            # for each prime factor, we can add it at any index i that we want
            pf = self.getPrimeFactors(num, primes)
            cur = 1
            for d in pf:
                ct = pf[d]
                v = n
                # there are (n + 1) choose k ways to add k prime factors
                for add in range(1, ct):
                    v *= (n + add)
                    v //= (add + 1)
                
                cur = (cur * v) % mod
                    
            ret = (ret + cur) % mod
                    
        return ret