[BOJ] 1715 카드 정렬하기

그리디, 우선순위 큐


해당 문제를 그리디로 접근한 정당성은 다음과 같다

  • 고르는 순서에 따라 결과가 달라진다. → 누적합을 구하는 문제인 만큼 매 순간 가장 작은 카드 두 개를 뽑아서 비교하면 될 것이라고 판단. → 매 순간 오름차순 정렬을 통해 가장 작은 두 개를 뽑으면 된다고 판단

처음에는 sort 한번만 하고 계속 더하는 누적 합 문제인줄 알았는데 틀리고 나서 손으로 다시 풀어보니 매 순간 오름차순 정렬을 해줘야 한다는 사실을 깨달았다.

python sort는 O(NlogN) 이걸 N번 한다면 O(N^2logN)이 되고, 이는 input의 범위가 1~100000인 해당 문제에서 시간 초과가 나올 것이 분명하다.

따라서, 자료구조 Min-heap을 통해서 구현하기로 결심 → Python library의 heapq 사용

import sys
import heapq

def solution(cards, n):  # 1 <= n <= 100000 -> O(NlogN) 까지 사용 가능
    if n == 1:
        return 0
    heap = list()
    # 해당 카드들을 모두 min-heap 에 저장
    for card in cards:              # O(N)
        heapq.heappush(heap, card)  # heappush: O(logN) 따라서 O(NlogN)
    result = list()
    while True:
        tmp = heapq.heappop(heap)   # 하나 뽑았을때
        if len(heap) == 0:          # 더이상 원소가 없다면 break
            break
        tmp += heapq.heappop(heap)  # 두 번째 뽑아서 더함
        heapq.heappush(heap, tmp)   # 더한 값 다시 min-heap 에 push
        result.append(tmp)          # 결과에도 push
    answer = sum(result)            # O(N)
    return answer

if __name__ == "__main__":
    N = int(sys.stdin.readline().rstrip())
    _cards = list()
    for _ in range(N):
        _cards.append(int(sys.stdin.readline().rstrip()))
    print(solution(_cards, N))