3. [Python] BOJ 2014 소수의 곱

1. 문제 링크

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

2014번: 소수의 곱

 

2. 문제 선정 과정 및 문제 내용

5월 과제 역시 자료구조를 활용한 도전적인 프로그래밍 문제 해결이었다. 나는 3, 4월에 공부했던 트리, 스택 말고 다른 자료구조를 사용하는 문제를 풀고 싶었고, 이에 선택한 문제가 바로 '힙' 자료구조를 이용하는 문제인 본 문제였다. 본 문제가 문제의 내용 자체를 이해하기에도 어렵지 않았고 내게 많은 흥미를 불러일으켰다. 또한 힙 자료구조를 사용하지 않고 반복문을 돌렸을 때와 힙 자료구조를 사용했을 때의 시간복잡도 차이도 궁금하여 본 문제를 선정하게 되었다.

 

문제 내용

문제:

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번째 오는 것을 출력한다.

 

3. 사용한 문제 풀이 전략

본 문제는 주어진 소수들을 활용하여 N번째로 작은 소수들의 곱을 시간 내에 출력하면 되는 문제였다. 따라서 소수의 곱들 중 최소값을 N번 찾아내는 과정을 빠르게 하는 것이 매우 중요했고, 이를 위해 '힙' 자료구조를 이용하였다.

 

3 - 1. 힙

 

*이진 트리: 한 노드가 최대 2개의 노드를 자식으로 가질 수 있는 트리

*완전 이진 트리: 이진 트리 중에서도 마지막 레벨을 제외하고 모든 레벨이 왼전히 채워져 있으며, 마지막 라벨의 모든 노드는 가능한 한 가장 왼쪽에 있는 트리. 완전 이진 트리에서 루트 노드의 번호를 0번이라고 하면, 현재 노드 번호를 i라 할 때 현재 노드의 부모 노드의 번호는  (i-1)//2이고, 현재 노드의 왼쪽 자식 노드의 번호는 i*2+1, 현재 노드의 오른쪽 자식 노드의 번호는 i*2+2이다.

 

힙은 완전 이진 트리 중 하나로, 부모 노드가 가진 자식 값은 자식 노드의 값보다 무조건 크거나(최대 힙) 무조건 작아야 한다(최소 힙). 또한 빈 노드가 없는 완전 이진 트리이므로 배열을 통해 구현할 수 있다. 이때 인덱스는 앞선 완전 이진 트리에 대한 설명과 같이 설정해주면 된다. 최대 힙과 최소 힙 각각의 구조를 살펴보면 다음과 같다.

 

최대 힙과 최소 힙의 구조

 

한편 힙은 우선순위 큐를 위하여 만들어진 자료구조이다. 우선순위 큐란 우선순위의 개념을 큐에 도입한 자료구조로, 기존 FIFO(First In, First, Out), 즉 먼저 들어온 자료가 먼저 나가는 선입선출의 큐 자료구조를 우선순위가 높은 데이터가 먼저 나가도록 설정한 자료구조이다. 우선순위 큐는 배열, 연결 리스트, 힙으로 구현 가능하지만 힙을 이용하는 것이 시간 복잡도가 O(logN)이 되어 가장 효율적이다.

 

힙에서 원소의 삽입이나 삭제가 일어나면 최대 힙의 조건이 깨질 수 있다. 이 경우 다시 힙의 조건을 만족하도록 노드의 위치를 바꾸는 것을 재구조화(heapify)라고 한다. heapify의 과정은 O(logN)의 시간복잡도를 가진다.

 

최소 힙의 경우를 예로 들어 힙 재구조화에 대하여 설명하겠다.

 

먼저, 새로운 원소가 최소 힙에 삽입되는 경우이다. 이때, 최소 힙의 가장 말단 노드(마지막 인덱스)에 원소를 삽입한다. 이후, 가장 말단 노드와 부모 노드의 값을 비교하면서 최소 힙 조건을 만족하는지 확인한다. 만약 부모 노드보다 삽입한 값이 더 작으면 부모 노드와 교환하고, 더 크면 정지하는 것이다. 아래는 위의 예시 최소 힙에 4라는 요소가 추가되는 경우를 예로 들어 설명한 그림이다.

최소 힙의 재구조화 과정 - 삽입

 

한편, 데이터 중 최소값을 찾기 위해 루트 노드를 삭제하게 되는 경우에도 재구조화가 필요하다. 이때는, 가장 말단의 원소를 루트 노드로 이동시킨 뒤, 루트 노드부터 자식 노드와 비교하여 최소 힙 조건을 만족하는지 확인한다. 부모 노드의 값보다 항상 자식 노드의 값이 크게 되어야 한다. 아래는 앞선 재구조화한 최소 힙에서 루트 노드가 삭제되는 경우의 재구조화를 나타낸 그림이다.

최소 힙의 재구조화 과정 - 루트 노드의 삭제

 

이렇게 힙과 힙의 재구조화 과정은 파이썬에서 최소 힙의 재구조화 과정에 대한 함수를 포함하는 모듈인 heapq의 heappush, heappop 함수를 사용하면 쉽게 구현할 수 있다. (최대 힙의 경우 적절히 부호를 변경하여 구현해야 한다) 그러나 위의 힙 재구조화 과정을 정확히 안다면 직접 힙 재구조화 과정을 파이썬 함수로 구현할 수 있다.

 

먼저, 최소 힙의 재구조화 과정 중 삽입의 경우(heappush)를 생각해보자. heap은 파이썬에서 리스트로 구현될 수 있으므로, heap이라는 기존 리스트와 삽입해야할 값 item이 매개변수로 들어오면 말단 노드에 item을 먼저 삽입한다. (heap.append(item)) 이후, 앞선 인덱스 관계식을 사용하여 부모 노드의 인덱스((i-1)//2)를 구하고, 만약 부모 노드가 item보다 더 작다면 작지 않게 될 때까지 부모 노드와 item을 교환해나가는 식으로 함수를 구성할 수 있다. 아래는 직접 구현한 코드이다.

def my_heappush(heap, item):
    heap.append(item)
    pos = len(heap) - 1
    while pos > 0:
        parpos = (pos - 1) // 2
        parent = heap[parpos]
        if parent <= item:
            break
        heap[pos], heap[parpos] = heap[parpos], heap[pos]
        pos = parpos

 

다음으로, 최소 힙의 재구조화 과정 중 루트 노드의 삭제의 경우(heappop)를 생각해보자. heappop 함수는 heap 리스트만을 매개 변수로 하고, 루트 노드를 삭제하면서 루트 노드를 반환한다. 따라서 return 값은 heap[0]이 되어야 한다. 한편 힙의 요소가 1개만 있지 않는 이상 마지막 요소를 가져와서 루트 노드로 올려야 하고 pop 매서드를 이용하면 쉽게 이를 코드로 구현할 수 있다. 또한 왼쪽 자식 노드 또는 오른쪽 자식 노드와 대소관계를 비교하여 최소 힙의 조건을 만족시켜야 하기 때문에(이때 인덱스는 마찬가지로 앞선 인덱스 관계식을 이용하여 구한다), 자식 노드가 없거나 (구한 자식 노드의 인덱스가 heap의 길이보다 크거나), 부모 노드가 두 자식 노드 중 더 작은 자식 노드보다 값이 작다면 교환을 멈추고, 그렇지 않다면 부모 노드와 가장 작은 자식 노드를 계속 교환해나가는 식으로 함수를 구현할 수 있다.

def my_heappop(heap):
    res = heap[0]
    last_element = heap.pop()
    if len(heap) > 0:
        heap[0] = last_element
        pos = 0
        while True:
            childpos1 = 2 * pos + 1
            childpos2 = 2 * pos + 2
            if childpos1 >= len(heap):
                break
            if childpos2 < len(heap) and heap[childpos2] < heap[childpos1]:
                min_child_pos = childpos2
            else:
                min_child_pos = childpos1
            if heap[min_child_pos] >= heap[pos]:
                break
            heap[pos], heap[min_child_pos] = heap[min_child_pos], heap[pos]
            pos = min_child_pos
    return res

 

 아래 링크에서 실제 heapq 모델에서 heappush, heappop 함수를 어떻게 구현하는지 그 코드를 볼 수 있다. 내가 구현한 것과 코드는 다르지만, 그 원리는 동일하다고 할 수 있다.

https://github.com/python/cpython/blob/3.12/Lib/heapq.py

cpython/Lib/heapq.py at 3.12 · python/cpython

 

3 - 2. 문제에 힙 자료구조 적용

 

본 문제는 힙 자료구조만 잘 이해하고 있다면 최소 힙을 이용해 비교적 쉽게 풀 수 있는 문제였다.

 

1. 최소 힙 정의 및 초기화: 문제에서 N, K를 입력받고 입력받은 K개의 소수들을 따로 리스트(nums)에 저장해놓은 뒤 my_heappush 함수를 이용하여 heap에 넣는다.

heap = []
K,N = map(int,input().split())
nums = list(map(int, input().split()))

for i in nums:
    my_heappush(heap, i)

 

2. heap 안의 최소 수(a)와 입력받은 소수들 중 하나/(nums[j])의 곱을 heap 안에 넣는다. 본 과정으로부터 heap 안에 입력받은 소수들 중 일부를 골라 곱했을 때의 모든 값이 들어가게 된다. 이때, 2x5와 5x2의 경우처럼 값이 중복될 수 있기 때문에, 이런 경우를 막기 위하여 a가 nums[j]로 나누어떨어지는지를 검사해야 한다. 만약 나누어떨어진다면, 이미 a에 nums[j]를 곱한 경우의 수를 전부 돌았다는 것이기 때문에 break문을 사용한다. 아래는 이를 구현한 코드이다.

a = my_heappop(heap)
for j in range(K):
    newnum = a*nums[j]
    my_heappush(heap, newnum)
    if a % nums[j] == 0:
        break

 

3. 정답 구하기: 문제에서 소수들의 곱 중 N번째로 작은 수를 구해야 하므로, a로 heap에서 가장 작은 수를 뽑아내는 과정을 N번 진행한다면 (for문을 N번 돌린다면) 최종적으로 남은 수(a)가 N번째로 작은 소수이다. 따라서 a를 출력하면 문제는 해결된다.

 

아래는 최종 코드이다. 전자는 heapq 모듈을 사용했을 경우, 후자는 사용하지 않았을 경우이다.

import heapq

heap = []
N,K = map(int,input().split())
nums = list(map(int, input().split()))

for i in nums:
    heapq.heappush(heap, i)

for i in range(K):
    a = heapq.heappop(heap)
    for j in range(N):
        newnum = a*nums[j]
        heapq.heappush(heap, newnum)
        if a % nums[j] == 0:
            break

print(a)

(heapq 모듈을 사용했을 경우)

def my_heappush(heap, item):
    heap.append(item)
    pos = len(heap) - 1
    while pos > 0:
        parpos = (pos - 1) // 2
        parent = heap[parpos]
        if parent <= item:
            break
        heap[pos], heap[parpos] = heap[parpos], heap[pos]
        pos = parpos

def my_heappop(heap):
    res = heap[0]
    last_element = heap.pop()
    if len(heap) > 0:
        heap[0] = last_element
        pos = 0
        while True:
            childpos1 = 2 * pos + 1
            childpos2 = 2 * pos + 2
            if childpos1 >= len(heap):
                break
            if childpos2 < len(heap) and heap[childpos2] < heap[childpos1]:
                min_child_pos = childpos2
            else:
                min_child_pos = childpos1
            if heap[min_child_pos] >= heap[pos]:
                break
            heap[pos], heap[min_child_pos] = heap[min_child_pos], heap[pos]
            pos = min_child_pos
    return res

heap = []
K,N = map(int,input().split())
nums = list(map(int, input().split()))

for i in nums:
    my_heappush(heap, i)

for i in range(N):
    a = my_heappop(heap)
    for j in range(K):
        newnum = a*nums[j]
        my_heappush(heap, newnum)
        if a % nums[j] == 0:
            break

print(a)

(모듈을 사용하지 않았을 경우)

 

4. 시행착오

본 문제는 처음에 힙 자료구조를 사용하지 않고 풀어보려고 했다. 그런데 그럴 경우 모든 소수의 곱의 경우의 수를 다 계산해본 뒤 최소값을 N번 걸러내는 매우매우 오래 걸리는 과정을 수행해야 했다. 그러나 이런 과정을 힙 자료구조를 이용해 구현한 우선순위 큐를 이용한다면 O(logN)의 상당히 빠른 시간 안에 해결할 수 있었다. 한편 또 고생했던 부분은 heappush, heappop 함수를 내가 직접 구현해본 부분이다. parent node와 child node에 대하여 검사하고 위치를 바꾸는 구문들을 실수 없이 작성하려니 상당히 까다로웠다. 실제로 제출해서 정답이 나왔어야 했기에, 모든 테스트케이스들을 통과할 수 있는 완벽한 코드를 구현해야 했는데 몇 번의 시행착오 끝에 결국 정답이 나왔을 때는 정말 기뻤다.

 

5. 느낀점

본 문제를 풀어보면서 힙과 우선순위 큐라는 자료구조에 대해 처음 알게 되었는데, 3월에 공부했던 세그먼트 트리와 마찬가지로 O(logN)의 시간 복잡도를 이용하여 선형 탐색을 할 경우 매우 오래 걸릴 수 있는 문제를 빠르게 풀어낼 수 있는 자료구조라 상당히 매력적이라고 생각했다. 이 힙 자료구조를 사용한다면 다른 최대값 또는 최소값을 구하는 문제들을 비교적 쉽게 시간 초과가 나지 않고 풀어낼 수 있을 것 같다.