[백준 2014번] 소수의 곱 (Python3)

from heapq import *

K,N = map(int,input().split())

prime = [*map(int,input().split())]

hq = [1]; nums = set()
MAX = 0
for _ in range(N+1):
  x = heappop(hq)
  for p in prime:
    if p*x in nums or len(hq)>N and p*x>MAX:
      continue
    MAX = max(MAX,p*x)
    heappush(hq,p*x); nums.add(p*x)
print(x)

우선순위 큐를 이용하여 해결하였다.

풀이 자체는 쉽게 떠올릴 수 있다. 하지만 숫자가 큰 탓인지 커팅을 하지 않고 그냥 무지성으로 코드를 짜면 메모리 초과가 발생한다. 따라서 이를 처리해주는 것이 필요하다.

 

풀이를 설명하면,

1. 가장 처음에는 힙큐에 1을 넣는다. nums는 중복방지를 위한 set이다.

2. 힙큐에서 가장 작은 수를 pop하고 이를 x로 둔다.

3. 모든 소수에 대해서 p*x를 힙큐에 넣는다. 이때 nums에 이미 들어가있는 숫자면 건너뛴다.

4. 만약 힙큐의 길이가 N보다 클 경우 지금까지의 모든 숫자 중 가장 큰 수보다 큰 경우 건너뛴다.

5. 2~4를 N+1번 실행하고, 가장 마지막 x를 출력한다.

 

4에서 힙큐의 길이가 N보다 크고, 지금까지의 모든 숫자보다 큰 경우에는 N번째 숫자가 될 가능성이 아예 없기 때문에 저장해줄 필요가 없게된다. 4번 과정을 하지 않으면 메모리 초과가 발생한다.

 

 

오늘의 교훈) 우선순위 큐를 적절하게 활용하자