Leetcode 313. Super Ugly Number

문제

A super ugly number is a positive integer whose prime factors are in the array primes.

Given an integer n and an array of integers primes, return the nth super ugly number.

The nth super ugly number is guaranteed to fit in a 32-bit signed integer.

• 가장 못생긴 숫자란, 주어진 prime 배열의 곱으로만 이루어진 양의 정수를 의미한다.
• 주어진 n 정수와 primes 소수 배열에 대해
• n번째로 큰 가장 못생긴 숫자 ( Super Ugly Number ) 를 구하시오

POINT

• 1은 primes가 [2, 3, 5] 일때 $1 = 2^0 * 3^0 * 5^0$ 이 성립한다.
• 즉 어떤 primes가 와도 1은 항상 super ugly number에 포함된다.

---

예시

• n = 12, primes = [2,7,13,19]

[1,2,4,7,8,13,14,16,19,26,28,32]

• 위 배열이 작은 ugly number부터 정확히 12개를 나열한 것이다.
• 즉 정답은 32 ( 12번째 )

---

제한

• $1 <= n <= 10^5$
• $1 <= primes.length <= 100$
• $2 <= primes[i] <= 1000$
• primes 배열의 정수는 반드시 소수이다.
• primes 배열은 오름차순 정렬되어 제공된다.

---

정답 코드

class Solution:
    def nthSuperUglyNumber(self, n: int, primes: List[int]) -> int:
        arr = [1]
        N = len(primes)
        indices = [0] * N

        h = [(v, i) for i, v in enumerate(primes)]
        heapify(h)

        while len(arr) < n:
            val, idx = heappop(h)

            if val != arr[-1]:
                arr.append(val)

            indices[idx] += 1
            next_val = primes[idx] * arr[indices[idx]]
            heappush(h, (next_val, idx))


        return arr[-1]

변수/구조 준비 해설

arr = [1]
N = len(primes)
indices = [0] * N

• arr: 지금까지 구한 정렬된 super ugly numbers를 담습니다. $\textbf{1}$로 시작합니다(모든 소수의 0제곱 곱이므로 항상 포함).

• indices: 길이 $N$의 배열로, 각 소수 primes[i]가 다음에 곱해야 할 arr의 위치를 가리킵니다. 즉, 다음 후보값은 항상

  $\text{candidate}_i = \text{primes}[i] \times \text{arr}[\text{indices}[i]].$

  시작은 모두 0이므로 초깃값은 primes[i] * arr[0] = primes[i] * 1 = primes[i]가 됩니다.

h = [(v, i) for i, v in enumerate(primes)]
heapify(h)

• h(min-heap): 튜플 (값, 소수의 인덱스)를 담습니다. 초기에는 각 소수 자체가 다음 후보이므로 (primes[i], i)를 모두 넣고 최소 힙으로 만듭니다.

메인 루프 동작

while len(arr) < n:
    val, idx = heappop(h)

    if val != arr[-1]:
        arr.append(val)

    indices[idx] += 1
    next_val = primes[idx] * arr[indices[idx]]
    heappush(h, (next_val, idx))

1. 최소 후보 뽑기: heappop(h)가 현재 가능한 후보들 중 가장 작은 값 val과 그것을 만든 소수 인덱스 idx를 줍니다.
2. 중복 제거 후 추가: 직전 추가 값(arr[-1])과 다르면 arr에 val을 추가합니다. (서로 다른 소수 경로가 같은 값을 만들 수 있으므로, 같은 값을 여러 번 넣지 않게 합니다.)
3. 해당 소수의 다음 후보 갱신: indices[idx]를 1 증가시켜, 같은 소수 primes[idx]를 다음 arr 원소에 곱해 새 후보 next_val = primes[idx] * arr[indices[idx]]를 만들고 힙에 다시 넣습니다.

이 과정을 arr 길이가 n이 될 때까지 반복하면, arr는 오름차순의 super ugly 수들을 앞에서부터 담게 됩니다.

작은 예시 추적

• n = 12, primes = [2, 7, 13, 19]

• 초기 상태

  • arr = [1]
  • indices = [0,0,0,0]
  • h = [(2,0), (7,1), (13,2), (19,3)]

반복 (핵심 단계만):

1. pop → (2,0) → arr[-1]=1과 다르므로 arr=[1,2]

   • indices[0]=1 → next=2*arr[1]=2*2=4 → push (4,0)

2. pop → (4,0) → arr=[1,2,4]

   • indices[0]=2 → next=2*arr[2]=2*4=8 → push (8,0)

3. pop → (7,1) → arr=[1,2,4,7]

   • indices[1]=1 → next=7*arr[1]=14 → push (14,1)

4. pop → (8,0) → arr=[1,2,4,7,8]

   • indices[0]=3 → next=2*arr[3]=14 → push (14,0)

5. pop → (13,2) → arr=[1,2,4,7,8,13]

   • indices[2]=1 → next=13*arr[1]=26 → push (26,2)

6. pop → (14,0) → arr[-1]=13과 다르므로 arr=[1,2,4,7,8,13,14]

   • indices[0]=4 → next=2*arr[4]=16 → push (16,0)

7. pop → (14,1) → 중복 값(이미 14 추가됨)이므로 arr에 추가하지 않음

   • indices[1]=2 → next=7*arr[2]=28 → push (28,1)

8. pop → (16,0) → arr=[1,2,4,7,8,13,14,16]

   • indices[0]=5 → next=2*arr[5]=26 → push (26,0)

9. pop → (19,3) → arr=[1,2,4,7,8,13,14,16,19]

   • indices[3]=1 → next=19*arr[1]=38 → push (38,3)

10. pop → (26,0) → arr=[1,2,4,7,8,13,14,16,19,26]

    • indices[0]=6 → next=2*arr[6]=28 → push (28,0)

11. pop → (26,2) → 중복(이미 26 있음) → 추가하지 않음

    • indices[2]=2 → next=13*arr[2]=52 → push (52,2)

12. pop → (28,0) → arr=[1,2,4,7,8,13,14,16,19,26,28]

    • indices[0]=7 → next=2*arr[7]=32 → push (32,0)

13. pop → (28,1) → 중복 → 추가하지 않음

    • indices[1]=3 → next=7*arr[3]=49 → push (49,1)

14. pop → (32,0) → arr=[1,2,4,7,8,13,14,16,19,26,28,32] → 길이 12 달성

따라서 12번째 값은 32입니다.

왜 빠짐없이, 정렬된 순서로 구해질까?

각 소수 $p_i$에 대해, 무한 수열 $S_i = \{p_i \times \text{arr}[t]\}_{t\ge0}$를 생각합니다. arr는 이미 정렬되어 있으므로 각 $S_i$도 오름차순입니다. 힙은 $S_0, S_1, \dots, S_{N-1}$의 현재 머리 원소들 중 최소를 뽑고, 그 수열의 다음 원소를 넣는 과정을 반복합니다. 이는 여러 정렬된 수열을 하나로 합치는 k-way merge와 동일하며, 따라서 결과 arr 역시 정렬되고 빠짐없이(중복은 제거) 생성됩니다.

중복은 서로 다른 i, j에 대해 $p_i \times \text{arr}[a] = p_j \times \text{arr}[b]$가 성립할 때 발생합니다. 우리는 if val != arr[-1]로 한 번만 삽입하고, 해당 값을 만든 각각의 수열에 대해 다음 후보로 이동하므로 합병이 올바르게 진행됩니다.

복잡도 분석

• 힙 크기는 항상 $N = \lvert\text{primes}\rvert\le 100$.
• 각 반복에서 heappop/heappush가 1회씩 → $\mathcal{O}(\log N)$ 비용.
• arr에 값이 실제로 추가되는 횟수는 정확히 $n$번. (중복 pop은 추가 없이 지나가지만, 각 중복도 곧 다음 후보로 넘어가므로 평균적으로 큰 부담이 되지 않음.)

시간복잡도: 대략 $\mathcal{O}(n\log N)$.
공간복잡도: arr가 $n$, 힙과 보조 배열이 $N$ → $\mathcal{O}(n+N)$.

자주 묻는 포인트 정리

• 왜 arr는 1로 시작하나요?
  어떤 소수 배열이 오더라도 $1 = \prod p_i^0$이므로 super ugly 수에 항상 포함됩니다. 시작을 1로 해야 이후 후보들(p_i * arr[idx])이 자연스럽게 소수 자체로 초기화됩니다.

• 중복 처리(if val != arr[-1])를 빼면?
  같은 값이 여러 번 들어가 arr가 잘못된 길이/내용이 됩니다. 반드시 필요합니다.

• 32-bit 보장과의 관계는?
  문제에서 $n$번째 값이 32-bit signed 범위에 든다고 보장하므로 파이썬에서는 오버플로우 걱정 없이 구현만 올바르면 됩니다.