[Python] 1715 카드 정렬하기

문제

정렬된 두 묶음의 숫자 카드가 있다고 하자. 각 묶음의 카드의 수를 A, B라 하면 보통 두 묶음을 합쳐서 하나로 만드는 데에는 A+B 번의 비교를 해야 한다. 이를테면, 20장의 숫자 카드 묶음과 30장의 숫자 카드 묶음을 합치려면 50번의 비교가 필요하다.

매우 많은 숫자 카드 묶음이 책상 위에 놓여 있다. 이들을 두 묶음씩 골라 서로 합쳐나간다면, 고르는 순서에 따라서 비교 횟수가 매우 달라진다. 예를 들어 10장, 20장, 40장의 묶음이 있다면 10장과 20장을 합친 뒤, 합친 30장 묶음과 40장을 합친다면 (10 + 20) + (30 + 40) = 100번의 비교가 필요하다. 그러나 10장과 40장을 합친 뒤, 합친 50장 묶음과 20장을 합친다면 (10 + 40) + (50 + 20) = 120 번의 비교가 필요하므로 덜 효율적인 방법이다.

N개의 숫자 카드 묶음의 각각의 크기가 주어질 때, 최소한 몇 번의 비교가 필요한지를 구하는 프로그램을 작성하시오.

입력

첫째 줄에 N이 주어진다. (1 ≤ N ≤ 100,000) 이어서 N개의 줄에 걸쳐 숫자 카드 묶음의 각각의 크기가 주어진다. 숫자 카드 묶음의 크기는 1,000보다 작거나 같은 양의 정수이다.

출력

첫째 줄에 최소 비교 횟수를 출력한다.

문제 풀이

A. 접근

정렬해서 풀면 될 것 같았는데 시간 복잡도 문제로 그렇게 풀어선 안 된다. 일단 입력이 10, 20, 40이라고 한다면, (10 + 20)을 먼저 해야 하는데, 그 이유는 더한 값을 다시 사용하기 때문이다. 즉, (10 + 40)을 한다면 나중에 (10 + 40)인 50을 또 계산에 포함해야 한다는 것이다. 그래서 (10 + 40) + (50 + 20) 으로 120이 나오는 것이다. 그렇기 때문에 매번 최솟값을 구해서 더해주어야 가장 최소한으로 카드를 정렬할 수 있다.

매번 최솟값, 즉 우선순위를 고려한다는 점이기 때문에 우선순위 큐 자료구조를 사용하도록 하자. 정렬을 계속해서 사용하면 시간 복잡도가 오버되기 때문에, 데이터의 추가와 삭제가 O(logN)이며 최솟값을 찾을 때는 O(1)인 우선순위 큐를 사용해야 한다.

B. 문제 풀이

전체 코드

import heapq
import sys
input = sys.stdin.readline

n = int(input())
q = list(int(input()) for _ in range(n))
temp = 0

heapq.heapify(q)

while len(q) != 1:
    a = heapq.heappop(q)
    b = heapq.heappop(q)
    heapq.heappush(q, a + b)
    temp += a + b

print(temp)

B-1. 준비

import heapq
import sys
input = sys.stdin.readline

n = int(input())
q = list(int(input()) for _ in range(n))
temp = 0

heapq.heapify(q)

먼저 우선순위 큐를 사용하기 때문에 heapq 라이브러리를 사용하며, sys 모듈로 입력을 받는다. temp는 비교를 담을 변수이고, q에서 각 리스트를 받은 다음에 heapq.heapify()를 통해서 최소 힙을 통한 우선순위 큐를 구현하도록 한다.

더보기

heapq.heappush(리스트, 넣을 값)으로 값을 추가하는 경우와 이미 존재하는 리스트를 heapq.heapify(리스트)로 구현하는 두 가지 경우가 있는데, 아무래도 리스트 컴프리헨션을 통해서 heapq.heapify()를 사용하는 것이 조금 더 빨랐다.

B-2. 풀이

while len(q) != 1:
    a = heapq.heappop(q)
    b = heapq.heappop(q)
    heapq.heappush(q, a + b)
    temp += a + b

print(temp)

큐에 남은 값이 한 개밖에 남지 않을 때까지 뽑고 넣고를 반복한다. 이때 비교한 수를 temp에 저장한다.

총평

우선순위 큐는 참 매력적인 자료구조인 것 같다. 넘 재밌당 ㅋㅅㅋ