Count the Number of Ideal Arrays LeetCode solution in Python

1
from math import sqrt
2

3

4
class Solution:
5
    def primesUpTo(self, n):
6
        primes = set(range(2, n + 1))
7
        for i in range(2, n):
8
            if i in primes:
9
                it = i * 2
10
                while it <= n:
11
                    if it in primes:
12
                        primes.remove(it)
13
                    it += i
14

15
        return primes
16

17
    def getPrimeFactors(self, n, primes):
18
        ret = {}
19
        sq = int(math.sqrt(n))
20

21
        for p in primes:
22
            if n in primes:
23
                ret[n] = 1
24
                break
25

26
            while n % p == 0:
27
                ret[p] = ret.get(p, 0) + 1
28
                n //= p
29

30
            if n <= 1:
31
                break
32

33
        return ret
34

35
    def idealArrays(self, n: int, maxValue: int) -> int:
36
        mod = 10**9 + 7
37
        ret = 0
38
        primes = self.primesUpTo(maxValue)
39

40
        for num in range(1, maxValue + 1):
41
            # find number of arrays that can end with num
42
            # for each prime factor, we can add it at any index i that we want
43
            pf = self.getPrimeFactors(num, primes)
44
            cur = 1
45
            for d in pf:
46
                ct = pf[d]
47
                v = n
48
                # there are (n + 1) choose k ways to add k prime factors
49
                for add in range(1, ct):
50
                    v *= n + add
51
                    v //= add + 1
52

53
                cur = (cur * v) % mod
54

55
            ret = (ret + cur) % mod
56

57
        return ret