[python] 2014 - 소수의 곱

2014번: 소수의 곱 (acmicpc.net)

2014번: 소수의 곱

 

우선순위 큐를 해결하는 문제였다.

처음 내가 생각한 풀이는 다음과 같았다.

 

1. 중복조합으로 소수의 곱을 계산하여 리스트에 넣는다.

2. 힙 정렬 후 heap을 n번 반복하여 n번째 위치한 값을 구한다.

 

그런데 착각한 것이 주어진 소수를 곱해나가면 무한히 가능하기 때문에 중복조합으로 계산을 하면 안됐다..!

심지어 중복조합 후 구한 값을 prod를 사용해 누적합을 구해서 계산했다.

 

결과는 맞게 출력되었지만 메모리 초과 문제가 발생하였다.

 

 

[수정 전]

import sys, heapq
from itertools import combinations_with_replacement
from math import prod

input = sys.stdin.readline
k,n = map(int, input().split())
lst=list(map(int, input().split()))
q = []
rst=0
for i in range(1, k+1):
    for combi in combinations_with_replacement(lst, i):
        q.append(prod(list(combi)))

heapq.heapify(q)

for _ in range(n):
    rst = heapq.heappop(q)
    
print(rst)

 

[수정 후]

import sys, heapq

input = sys.stdin.readline

k,n = map(int, input().split())
lst=list(map(int, input().split()))

heap=[]
visited=set()
max_value=max(lst)

for i in lst:
    heapq.heappush(heap, i)
    
for i in range(n-1):
    ele = heapq.heappop(heap)
    
    for x in lst:
        num = ele*x
        
        if len(heap) >=n and max_value <num:
            continue
        if num not in visited:
            heapq.heappush(heap, num)
            max_value = max(max_value, num)
            visited.add(num)
            
            
print(heapq.heappop(heap))

 

우선 기존 소수 리스트를 힙에 넣고 최소값을 뽑아내 리스트의 원소들과 곱하여 리스트에 추가하는 과정을 반복한다.

n-1번 반복하는 이유는 n번째 위치하는 수를 구하기 위함이다.

출력할 때 heappop을 해주었기 때문에 n-1만 반복하였다.

 

리스트의 최소값과 기존 리스트 원소를 곱한 값이 힙의 길이보다 같거나 크면서 리스트의 최대값보다 크다면

이 값은 뽑아낼 일이 없기 때문에 무시하도록 하였다.