[python] 백준 2014 :: 소수의 곱 (우선순위 큐)

[소수의 곱]

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

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

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

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

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

풀이

01. 소수를 힙에 넣는다.

• 입력받은 소수를 최소힙에 index와 함께 넣는다.

• 힙의 pop을 이용해 작은 수부터 하나씩 꺼내올 예정!

heap = [] # 최소힙

for idx, p in enumerate(prime): # 주어진 소수들을 힙에 추가한다.
    heapq.heappush(heap,(p, idx)) # 값, 초기 소수의 인덱스

02. 힙에서 하나씩 꺼내서 다른 소수들과 곱한다.

• n번째 숫자가 필요하니까 n번 반복한다.

for i in range(n): # 힙의 n번째 원소 출력

• 힙에서 제일 작은 값과 함께 넣은 index를 꺼낸다.

    # 힙에서 제일 작은 값과 그 값을 만들어낸 초기 소수의 인덱스를 꺼낸다.
    min_number, base_idx = heapq.heappop(heap)

• 여기에 소수를 한번씩 다 곱해준다.

• 이렇게 만들어진 수를 힙에 추가할 때도,
  처음에 넣을 때와 마찬가지로 곱한 소수의 인덱스를 같이 넣어준다.

• 🚨 이렇게 그냥 곱해서 넣으면, 중복이 발생한다 😱
  중복 제거는 아래에서 설명!

    for idx, p in enumerate(prime):

        # 힙에서 꺼낸 제일 작은 원소에 소수를 차례대로 하나씩 곱한다.
        new_number = p * min_number
        
        # 만든 새 숫자를 힙에 추가한다.
        heapq.heappush(heap, (new_number, idx))

03. 중복 제거를 위한 조건 추가

• 아래의 조건을 for문 시작 부분에 넣어서, 이 경우에 해당하면 숫자를 더 곱하지 않도록 해야 한다.

• 이 조건을 추가하지 않은 경우에 연산되는 결과를 모두 출력해보면 아래처럼 나온다.

꺼낸 숫자(min_number):  2 꺼낸 숫자의 인덱스(base_idx):  0
곱한 소수(p):  2 곱한 소수의 인덱스(idx):  0
새로 만든 숫자(new_number):  4
###################################
꺼낸 숫자(min_number):  2 꺼낸 숫자의 인덱스(base_idx):  0
곱한 소수(p):  3 곱한 소수의 인덱스(idx):  1
새로 만든 숫자(new_number):  6
###################################
꺼낸 숫자(min_number):  2 꺼낸 숫자의 인덱스(base_idx):  0
곱한 소수(p):  5 곱한 소수의 인덱스(idx):  2
새로 만든 숫자(new_number):  10
###################################
꺼낸 숫자(min_number):  2 꺼낸 숫자의 인덱스(base_idx):  0
곱한 소수(p):  7 곱한 소수의 인덱스(idx):  3
새로 만든 숫자(new_number):  14
###################################
꺼낸 숫자(min_number):  3 꺼낸 숫자의 인덱스(base_idx):  1
곱한 소수(p):  2 곱한 소수의 인덱스(idx):  0
새로 만든 숫자(new_number):  6

...

• 이때, 2를 꺼내서 3을 곱한 것과 3을 꺼내서 2를 곱한 것은 동일하므로 한번만 수행해야 한다.
  이걸 막아주기 위해서 heap에 저장할 때 인덱스를 함께 저장한 것이다.
  꺼낸 숫자의 인덱스가 소수의 인덱스보다 작거나 같을때까지만 연산을 수행해주면 이런 반복을 제거할 수 있다.

        # 초기 소수의 다음 소수부터는 곱하지 않는다. (중복 방지)
        if idx > base_idx:
            break

끝~!😘

전체 코드

import heapq
k, n = map(int, input().split())
prime = list(map(int, input().split()))
heap = [] # 최소힙

for idx, p in enumerate(prime): # 주어진 소수들을 힙에 추가한다.
    heapq.heappush(heap,(p, idx)) # 값, 초기 소수의 인덱스


for i in range(n): # 힙의 n번째 원소 출력

    # 힙에서 제일 작은 값과 그 값을 만들어낸 초기 소수의 인덱스를 꺼낸다.
    min_number, base_idx = heapq.heappop(heap)

    for idx, p in enumerate(prime):
        
        # 초기 소수의 다음 소수부터는 곱하지 않는다. (중복 방지)
        if idx > base_idx:
            break

        # 힙에서 꺼낸 제일 작은 원소에 소수를 차례대로 하나씩 곱한다.
        new_number = p * min_number
        
        # 만든 새 숫자를 힙에 추가한다.
        heapq.heappush(heap, (new_number, idx))


    if i == n-1:
        print(min_number)