[TIL] 22.08.13 SAT

정렬 알고리즘

카드 정렬하기

정렬 카테고리에 있는 문제라서 처음에는 여러가지 정렬을 사용하여 접근해보았다. 문제를 해결하는 로직은 크기가 작은 두 개의 카드 뭉치를 더해주고, 합해진 카드 뭉치를 다시 모든 카드 뭉치와 비교하여 작은 것들끼리 합쳐주면 된다.

그렇게 처음 작성한 코드가 아래와 같다.

처음 작성한 코드

n = int(input())

deck = []

for _ in range(n):
	deck.append(int(input())

_sum = 0

if len(deck) == 1:
	_sum = 0
    
deck.sort(reverse=True)

while len(deck) > 1:
	a = deck.pop()
	b = deck.pop()
	c = a + b
	_sum = _sum + c

	deck.append(c)
    deck.sort(reverse=True)

print(_sum)

전체 카드뭉치가 담겨 있는 배열을 내림차순으로 정렬하고, 가장 작은 두 개의 카드 뭉치를 가져온다. 두 개의 카드 뭉치를 더하여 전체 비교 횟수에 추가해주고, 합쳐진 카드 뭉치는 deck 배열에 삽입하여 다시 한 번 역순으로 정렬한다. 이를 deck 배열에 남은 카드 뭉치가 없을 때까지 반복한다.

하지만 이렇게 했더니 시간 초과 오류가 발생한다. 그도 그럴 것이 카드 뭉치의 개수 n의 크기가 최대 100,000개이다. 결국 최악의 경우 100,000!만큼 비교를 해야 하기 때문에 시간 초과 오류가 발생하는 게 당연했다.

그래서 생각해낸 게 heap 자료구조를 사용하는 것이다.

heap은 우선순위 큐를 구현하기 위해 사용되는 자료구조로, 배열에 들어간 순서에 관계 없이 우선순위가 가장 높은 데이터에 먼저 접근하는 것을 말한다.

그렇게 작성한 코드는 다음과 같다.

힙 자료구조를 이용해 작성한 코드

import heapq

n = int(input())

deck = []

for _ in range(n):
	heapq.heappush(deck, int(input()))

_sum = 0

if len(deck) == 1:
	_sum = 0

while len(deck) > 1:
	a = heapq.heappop(deck)
	b = heapq.heappop(deck)
	c = a + b
	_sum = _sum + c

	heapq.heappush(deck, c)

print(_sum)

heapq.heappop()을 했을 때, 우선순위가 높은 데이터를 먼저 반환하므로 역순으로 정렬하는 과정이 필요가 없고, 그렇게 때문에 이전에 작성한 코드보다 훨씬 빠른 시간 복잡도를 갖는다.

--

참고 도서

🙇🏻‍♂️ 나 동빈, 이것이 취업을 위한 코딩테스트다