[백준(BOJ)] #2014- 소수의 곱 (파이썬, PyPy3)

문제

https://www.acmicpc.net/problem/2014

2014번: 소수의 곱

 

---

문제 풀이

우선순위 큐를 사용해 문제를 해결할 수 있습니다.

 

2022.07.19 - [CS/자료구조] - [파이썬으로 배우는 자료구조] 우선순위 큐(Priority Queue), 힙(Heap)

[파이썬으로 배우는 자료구조] 우선순위 큐(Priority Queue), 힙(Heap)

 

기존 풀이

 

1) 입력받은 소수 리스트를 a에 저장하고, 소수 리스트로 구성된 heap 준비
2) heap에서 pop을 하여 t를 얻음
3) t가 이미 처리된 숫자이면 4) 과정 생략 (카운트 세지 않음)
4) a에 저장된 소수들을 순회하면서 t를 곱하여 얻은 수들을 heap에 저장
5) 2) ~ 4) 과정을 N번 반복하고 마지막에 pop 한 숫자가 N번째 오는 숫자가 됨

입력받은 소수들과 우선순위 큐에 저장된 숫자들을 모두 곱하여 얻을 수 있는 숫자들을 순서대로 구하고, N번 반복하여 N번째 오는 숫자를 구하는 것입니다. 

 

세 소수 2, 5, 7을 입력받아 7번째 숫자를 구하는 과정을 함께 봐보겠습니다.

 

 

1) 힙 [2, 5, 7]에서 pop -> 1번째 숫자 = 2 -> 2 x [2, 5, 7] = [4, 10, 14] -> 힙에 push -> 힙 = [4, 7, 5, 10, 14]

 

2) 힙 [4, 7, 5, 10, 14]에서 pop -> 2번째 숫자 = 4 -> 4 x [2, 5, 7] = [8, 20, 28] -> 힙에 push -> 힙 = [5, 7, 14, 10, 8, 20, 28]

 

3) 힙 [5, 7, 14, 10, 8, 20, 28]에서 pop -> 3번째 숫자 = 5 -> 5 x [2, 5, 7] = [10, 25, 35] -> 힙에 push -> 힙 = [7, 8, 10, 10, 28, 20, 14, 25, 35]

 

4) 힙 [7, 8, 10, 10, 28, 20, 14, 25, 35]에서 pop -> 4번째 숫자 = 7 -> 7 x [2, 5, 7] = [14, 35, 49] -> 힙에 push -> 힙 = [8, 10, 10, 14, 28, 20, 14, 35, 25, 35, 49]

 

5) 힙 [8, 10, 10, 14, 28, 20, 14, 35, 25, 35, 49]에서 pop -> 5번째 숫자 = 8 -> 8 x [2, 5, 7] = [16, 40, 56] -> 힙에 push -> 힙 = [10, 10, 14, 14, 16, 20, 49, 35, 25, 35, 28, 40, 56]

 

6) 힙 [10, 10, 14, 14, 16, 20, 49, 35, 25, 35, 28, 40, 56]에서 pop -> 6번째 숫자 = 10 -> 10 x [2, 5, 7] = [20, 50, 70] -> 힙에 push -> 힙 = [10, 14, 14, 25, 16, 20, 49, 35, 56, 35, 28, 40, 20, 50, 70]

 

7) 힙 [10, 14, 14, 25, 16, 20, 49, 35, 56, 35, 28, 40, 20, 50, 70]에서 pop -> 10은 이미 처리된 숫자이므로 생략 -> 힙 = [14, 14, 20, 25, 16, 20, 49, 35, 56, 35, 28, 40, 70, 50]

 

8) 힙 [14, 14, 20, 25, 16, 20, 49, 35, 56, 35, 28, 40, 70, 50]에서 pop -> 7번째 숫자 = 14 -> 14 x [2, 5, 7] = [28, 70, 98] -> 힙에 push -> 힙 = [14, 16, 20, 25, 28, 20, 28, 35, 56, 35, 50, 40, 70, 49, 70, 98]

 

9) 7번째 숫자는 14

 

개선된 풀이

저의 원래 풀이는 결과는 맞지만 비효율적인 풀이입니다. 무작정 힙에 삽입을 하고 힙에서 pop을 하면서 중복 검사를 하는 로직이기 때문이죠. 따라서 다른 풀이들을 참고하여 중복 검사를 미리하고, 중복되는 숫자를 힙에 넣지 않는 방법으로 풀이를 개선했습니다.

 

1) 입력받은 소수 리스트를 a에 저장하고, 소수 리스트로 구성된 heap 준비
2) heap에서 pop을 하여 t를 얻음
3) a에 저장된 소수들을 순회하면서 t를 곱하여 얻은 수들을 heap에 저장
4) heap에 저장하고 다음 소수를 순회하기 전에 t가 현재 소수로 나누어 떨어지면 break
5) 2)~4) 과정을 N번 반복하고 마지막에 pop 한 숫자가 N번째 오는 숫자가 됨

 

 

pop 해서 얻은 숫자 t를 a의 모든 소수와 곱한 값들을 전부 heap에 저장하지 않고, t가 소수와 나누어 떨어지면 더 이상 heap에 저장하지 않습니다. 따라서 위와 같이 heap에 저장된 수가 확연히 적다는 것을 확인할 수 있습니다.

 

---

구현

기존 코드

import sys
import heapq as hq
input = sys.stdin.readline

if __name__ == "__main__":
    K, N = map(int, input().split())
    a = list(map(int, input().split()))
    heap = [x for x in a]

    cnt = 0
    res = 0

    while cnt < N:
        t = hq.heappop(heap)
        # pop한 값이 이미 처리되지 않았을 경우만
        if t != res:
            res = t
            cnt += 1
            for ai in a:
                hq.heappush(heap, ai * t)

    print(res)

기존 코드 실행 결과

 

개선된 코드

import sys
import heapq as hq
input = sys.stdin.readline

if __name__ == "__main__":
    K, N = map(int, input().split())
    a = list(map(int, input().split()))
    heap = [x for x in a]

    res = 0

    for i in range(N):
        res = hq.heappop(heap)
        for ai in a:
            hq.heappush(heap, ai * res)
            # 중복 검사
            if res % ai == 0:
                break
    print(res)

개선된 코드 실행 결과

메모리와 시간이 확연히 차이 나는 것을 확인할 수 있습니다.