1715 - 카드 정렬하기

info

문제 보기: 1715 - 카드 정렬하기
소요 시간: 17분
풀이 언어: python
체감 난이도: 2️⃣
리뷰 횟수: ✅

풀이 키워드

스포주의

자료구조 힙

풀이 코드

info

메모리: 36292 KB
시간: 124 ms

import sys
import heapq
input = sys.stdin.readline

n = 0
heap = []

def solution():
    global n, heap
    n = int(input())
    for i in range(n):
        heapq.heappush(heap, int(input()))
    
    ans = 0
    
    # edge case
    if n == 1:
        print(0) # no need to compare
        return
    
    # else
    while True:
        lhs = heapq.heappop(heap)
        try:
            rhs = heapq.heappop(heap)
        except:
            break
        
        # merge
        k = lhs + rhs
        heapq.heappush(heap, k)
        ans += k
        
    print(ans)
    

if __name__ == '__main__':
    solution()

풀이 해설

정렬되어있지 않은 데이터 리스트를 매 번 최소의 비용으로 정렬하는 것이 핵심인 문제이다.

비교 횟수 계산을 누적합으로 하되, 최소값이 되어야 하므로 작은 값부터 합해야 하는 점이 point이다.

하지만 한 번 값을 합치면 나머지 데이터 리스트에서 이 합이 몇 번째 큰 값인지 선형시간에 알 수 없으므로,

이 경우에 어울리는 최선의 자료구조로 heap을 생각해볼 수 있다.

📌파이썬의 `heapq`

heap 정렬의 경우 일반적으로 $O(logn)$ 이다.

파이썬 heapq의 경우 heappush와 heappop 모두 $O(logn)$ 의 비용을 가지고 있다.

📌Corner case

$1 ≤ N ≤ 100{,}000$ 이므로 n은 1이 될 수도 있다.

이 경우 시작부터 비교 횟수를 계산할 필요가 없으므로 0을 return 한다.

if n == 1:
    print(0) # no need to compare
    return

📌Time complexity

for i in range(n):
    heapq.heappush(heap, int(input()))

초반에서,
heap에 데이터를 넣어 정렬하는 부분은 $O(logn)$ 비용이 드는 heappush를 n번 수행하므로, $O(nlogn)$ 이다.

while True:
    lhs = heapq.heappop(heap)
    try:
        rhs = heapq.heappop(heap)
    except:
        break
    
    # merge
    k = lhs + rhs
    heapq.heappush(heap, k)
    ans += k

그 후 비교 횟수를 계산하는 while문에서, 2번의 heappop과 1번의 heappush가 매 iteration마다 일어나고 있다.

또한 iteration은 총 n만큼 발생하므로, 복잡도는 $T(3nlogn) == O(nlogn)$ 이다.

따라서 풀이코드의 최종 시간복잡도는 $O(nlogn)$ 이다.

메모

heap 내부의 원소를 합치다보면 무조건적으로 결국엔 1개의 원소만 남게 된다. 이 경우 rhs를 구하는 statement에서 RTE가 터지게 된다. 따라서 해당 부분을 try-except 트릭을 활용하여 처리해보았다.
PriorityQueue 모듈도 있는데, 동기화(Thread-safe) 기능 때문에 heapq보다 느리다고 한다.

풀이 키워드​

풀이 코드​

풀이 해설​

📌파이썬의 heapq​

📌Corner case​

📌Time complexity​

메모​