JUNGLE_2주차

컴퓨팅 사고로의 전환 1주 차

 

주차별 키워드

- 이분 탐색, 분할 정복, 스택, 큐, 우선순위 큐

 

더 고민해볼 키워드

- 세그먼트 트리

- 분할 정복

- heapq, stack, deque Class 구현

---

 

SW사관학교 정글 WEEK01 시험

- 쿼드 트리 click

- 쇠막대기 click

- 소수의 곱 click

 

 

소수의 곱 click

2014번: 소수의 곱

 

이번 주의 마지막 문제였던 숫자 야구에 대해서 알아보자.

 

K 개의 소수가 주어졌을 때, 이 소수들을 서로 곱해 오름차순으로 나타낸다. 이때, N번째에 위치하는 수를 출력하는 문제이다.

 

카드 정렬하기 문제와 비슷하게 최소 값들을 heap에서 뽑아내며 이들을 K개의 소수와 곱한 수들을 다시 heap에 추가해주며 풀어나가면 된다.

 

하지만, 문제점이 두 가지 있다.

 

1. 중복

첫 번째는 최솟값을 뽑고 K개의 소수와 곱해서 추가하는 과정을 반복하다 보면, 중복되는 수가 생긴다는 것이다.

예제를 활용해서 이해해보자.

 

• 2, 5, 7 이렇게 세 소수가 있을 때, 2를 뽑아서 2, 5, 7과 곱해서 4, 10, 14를 추가한다.
  • 첫 번째 숫자는 2, heap = [4, 5, 7, 10, 14]
• 4를 뽑아서 2, 5, 7과 곱한 후에 추가해주자.
  • 두 번째 숫자는 4, heap = [5, 7, 8, 10, 14, 20, 28]
• 5를 뽑아서 2, 5, 7과 곱한 후에 추가해주자. 10, 25, 35가 추가된다.
  
  • 그런데 10은 이미 첫 번째 단계에서 추가되어 있는 숫자이다. 

이렇게, 단순히 heap에서 최솟값을 뽑고 주어진 소수와 곱해서 heap에 추가하면, 2 * 5 = 10, 5 * 2 = 10처럼 중복이 발생하게 된다. 

 

Heap에 index와 함께 추가 

 

Heap에 [2, 5, 7]과 같이 단순히 숫자를 추가하는 것이 아닌, 곱한 숫자들의 인덱스를 함께 추가해 준다.

 

• [(2, 0), (5, 1), (7, 2)]
  • (2, 0)을 POP, index가 0이기 때문에 인덱스가 0 이상인 숫자들만 곱해서 PUSH
  • (2*2, 0), (2*5, 1), (2*7, 2)을 PUSH
• [(4, 0), (5, 1), (7, 2), (10, 1), (14, 2)]
  
  • (4, 0)을 POP, index가 0이기 때문에 인덱스가 0 이상인 숫자들만 곱해서 PUSH
  • (4*2, 0), (4*5, 1), (4*7, 2)을 PUSH
• [(5, 1), (7, 2), (8, 0), (10, 1), (14, 2), (20, 1), (28, 2)]
  • (5, 1)을 POP, index가 1이기 때문에 인덱스가 1 이상인 숫자들만 곱해서 PUSH
  • (5*5, 1), (5*7, 2) PUSH

....

이렇게 중복을 피해 추가할 수 있다.

 

 

2. 메모리 초과

 

두 번째는 K개(최대 100개)의 소수들과 곱해서 추가하는 과정을 N번(최대 10만 회) 반복하면, N번째로 큰 수 뒤로 굉장히 많은 수가 따라붙어 메모리 초과가 발생한다는 것이다.

 

마지막에 heapq에 어떤 수가 추가되어있는지 출력해 보았다.

예제 입력을 넣었을 때, heapq에 어떤 수가 포함되어있는지 출력해 보았다. 19번째 숫자는 27이지만, 147까지 35개의 숫자가 추가로 포함되어 있음을 알 수 있다. 이렇게 N번째 숫자 뒤로 추가되는, 답을 도출하는데 필요 없는 숫자들을 제거해 주어야 했다.

 

기존의 Heap과 별도로 기존에 추가된 값들 중 최댓값을 저장하는 Heap(max_num)을 한 개 더 생성했다.

 

max_num의 길이가 N이상일 때, max_num[0](추가된 N개 이상의 수들 중 최댓값) 보다 큰 수들은 정답이 될 가능성이 없는 수다.

 

따라서, 추가하기 전 max_num(최대 힙)에 저장되어 있는 최댓값과 비교하고, 더 작을 때만 추가해도 N번째 최솟값을 찾을 수 있다.

 

 

Code

import sys
read=sys.stdin.readline
import heapq as hq
    
K, N = map(int,read().split())
prime = list(map(int,read().split()))
if K == 1:
    print(prime[0]**N)
    exit()

num=list((prime[i], i) for i in range(K))
max_num = list(-prime[i] for i in range(K))
hq.heapify(num)
hq.heapify(max_num)

for x in range(N):
    num_pop, index = hq.heappop(num)
    for i in range(index, K):
        if len(max_num) >= N and -max_num[0] <= num_pop*prime[i]:
            continue
        hq.heappush(num, (num_pop*prime[i], i))
        hq.heappush(max_num, -num_pop*prime[i])

print(num_pop)

 

틀린 부분이 있을 수 있습니다. 피드백 주시면 고치도록 하겠습니다! 감사합니다.👍

[꼭 다시 풀어보기]