Problem link
https://leetcode.com/problems/count-ways-to-make-array-with-product/
Problem Summary
n, k 가 입력으로 주어질 때 n개의 수의 곱으로 k를 만들 수 있는 방법의 수를 구하는 문제.
Solution
처음에 단순 DP로 하니 시간 초과가 난다... O(NKD) (F: 약수 개수)
dicussion을 보니 Stars and bars 개념을 이용해서 풀 수 있다.
Stars and bars를 간단하게 설명하면 별을 바로 나눌 수 있는 경우의 수를 구하는 문제이다. 즉, 별과 바를 일렬로 배치하는 경우의 수라고 보면 된다.
이해하기 쉬운 Numberphile 유튜브 영상
예시로 n=3, k=16의 경우 2x2x2x2 이므로 별이 4개이고 n=3이므로 바는 2개라고 볼 수 있다. 그럴 때 나누는 경우의 수는
(4 + 3 - 1) C (3 - 1) = 6C2
가 된다.
단, stars and bars는 같은 수에만 적용이 되므로 이 문제에서는 소인수분해를 해서 같은 소수끼리의 경우의 수를 구하고 다 곱해주면 된다.
또 다른 예시로 n=3, k=12 (2 x 2 x 3)을 보면
1. 2*2 를 3개로 채우는 경우의 수 C(4, 2) = 6
{ ||2 2 -> [], [], [4] }, { |2|2 } -> [], [2], [2] }, { |2 2| } -> [], [4], [] },
{ 2||2 } -> [2], [], [2] }, { 2|2| } -> [2], [2], [] }, { 2 2|| } -> [4], [], [] }
2. 3을 3개로 채우는 경우의 수 C(3, 1) = 3
{ ||3 -> [], [], [3]}, { |3| -> [], [3], [] }, { 3|| -> [3], [], [] }
2, 3을 독립적으로 보고 두 경우의 수를 곱하면 6*3=18이 된다.
참고: https://leetcode.com/problems/count-ways-to-make-array-with-product/discuss/1035607/C%2B%2BPython-Precompute
https://leetcode.com/problems/count-ways-to-make-array-with-product/discuss/1032567/Python-100-Primes-factorization-and-combinations-%2B-well-explained-%2B-thinking-analysis
Source Code
from functools import lru_cache
from math import comb
from typing import List
class Solution:
def waysToFillArray(self, queries: List[List[int]]) -> List[int]:
primes = []
for i in range(2, 100):
is_prime = True
for j in range(2, i):
if i % j == 0:
is_prime = False
break
if is_prime:
primes.append(i)
def count(nn, kk):
ret = 1
for p in primes:
cnt = 0
while kk % p == 0:
kk //= p
cnt += 1
ret *= comb(nn + cnt - 1, cnt)
if kk > 1:
ret *= comb(nn, 1)
return ret % (10 ** 9 + 7)
res = []
for n, k in queries:
res.append(count(n, k))
return res