[백준/Baekjoon] 2014번: 소수의 곱

1. 문제

K개의 소수가 있다. 이때, 이 소수들 중에서 몇 개를 곱해서 얻게 되는 수들이 있을 것이다. 소수들을 선택할 때에는 같은 수를 선택해도 되며, 주어지는 소수 자체도 포함시키자.

예를 들어 세 소수가 2, 5, 7이었다면, 이러한 곱들을 오름차순으로 나타내 보면, 2, 4, 5, 7, 8, 10, 14, 16, 20, 25, 28, 32, 35, 등이 된다.

K개의 소수가 주어졌을 때, 이러한 소수의 곱들 중에서 N번째 수를 구해 보자. 단 정답은 231보다 작은 자연수이다.

2. 입력

첫째 줄에 K(1 ≤ K ≤ 100), N(1 ≤ N ≤ 100,000)이 주어진다. 다음 줄에는 K개의 소수가 오름차순으로 주어진다. 같은 소수가 여러 번 주어지는 경우는 없으며, 주어지는 소수는 모두 541보다 작거나 같은 자연수이다.

3. 출력

첫째 줄에 문제에서 설명한 대로 소수의 곱을 나열했을 때, N번째 오는 것을 출력한다.

4. 예제 입력 1

4 19
2 3 5 7

5. 예제 출력 1

27

6. 풀이

힙(Heap) 은 힙의 특성(최소 힙(Min Heap) 에서는 부모가 항상 자식보다 작거나 같다) 을 만족하는 거의 완전한 트리인 특수한 트리 기반의 자료구조다. 이를 구현하기 위해 heapq 모듈 을 사용한다.

소수의 곱을 나열하는 방법은 다음과 같다.

1. 입력받은 소수들을 힙에 삽입한다. (이 때 소수들을 저장하는 목록은 따로 존재한다.)
2. 힙에서 가장 작은 값을 꺼내 소수들과 곱하여 값을 힙에 삽입한다.
3. 2번 과정을 N번 반복한다. 이 때 N번째 꺼내는 값이 소수의 곱에서 N번째 값이다.

for _ in range(n):
    num = heapq.heappop(heap)  # 힙에서 가장 작은 값을 꺼냄
    for i in range(k):
        temp = num * data[i]  # 힙에서 꺼낸 값을 소수들과 곱함
        heapq.heappush(heap, temp)  # 곱한 값을 힙에 삽입

이러한 과정을 그대로 구현하면 문제가 발생한다. 이는 중복되는 소수의 곱이 힙 안에 존재하게 되기 때문이다. 예를 들어보면, 세 소수가 2, 5, 7인 경우, 2 * 5, 5 * 2와 같이 중복이 존재할 수 있다. 이러한 중복은 힙에서 꺼낸 값을 입력받은 소수들로 나누어보는 방법으로 해결할 수 있다. 입력받은 소수로 나누어 떨어지는 경우까지만 힙에 삽입하면 중복 없이 소수의 곱을 구할 수 있다.

if num % data[i] == 0:  # 입력받은 소수로 나누어 떨어지는 경우
    break  # 이 경우까지만 힙에 삽입

코딩중독 님의 블로그 에 표와 함께 더 자세한 설명으로 이루어져 있다.

7. 코드

import heapq

k, n = map(int, input().split())
data = list(map(int, input().split()))
heap = []

for d in data:
    heapq.heappush(heap, d)

for _ in range(n):
    num = heapq.heappop(heap)
    for i in range(k):
        temp = num * data[i]
        heapq.heappush(heap, temp)

        if num % data[i] == 0:
            break

print(num)