1735. Count Ways to Make Array With Product

[zeikar]

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 유튜브 영상

<iframe width="560" height="315" src="https://www.youtube.com/embed/UTCScjoPymA" title="YouTube video player" frameborder="0" allow="accelerometer; autoplay; clipboard-write; encrypted-media; gyroscope; picture-in-picture" allowfullscreen></iframe>

예시로 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