[알고리즘/Python] 항상 작은 숫자 2개를 활용하고 싶을 때 (list, heapq)

1. 문제 상황 및 해결 방법

• 주어진 문제는 리스트에서 가장 작은 두 숫자를 반복적으로 찾아 합치고, 이를 다시 리스트에 추가하여 최종 비용을 계산하는 것입니다.
• 기존 코드에서는 리스트를 반복적으로 정렬하여 가장 작은 두 숫자를 찾고 제거합니다.
• 이 코드로 문제에 제출하니 시간 초과가 떠서 해결방안을 고민했습니다.

n = int(input())
nums = list(map(int, input().split()))
cost = 0
while (len(nums) > 1) :
    # 리스트 정렬하여 가장 작은 두 숫자 구하기
    nums.sort()
    n1 = nums.pop(0)
    n2 = nums.pop(0)

    # 드는 비용 업데이트
    cost += (n1+n2)
    
    # 두 숫자를 리스트에서 빼서 더하고, 맨 뒤의 리스트에 다시 넣어주기
    nums.append(n1+n2)

print(cost)

 

2. 해결방안 : 힙큐(heapq)를 사용

 

heapq: 알아서 작은 순서대로 정렬해주는 큐

• heapq는 최소 힙(min-heap)을 제공하여, 가장 작은 요소를 빠르게 찾고 제거할 수 있는 자료구조입니다.

시간 초과 문제

• 리스트를 반복적으로 정렬하는 작업은 매우 비효율적입니다.
• Python의 sort() 함수는 퀵 정렬 알고리즘을 사용하며, 평균 시간 복잡도는 𝑂(𝑛log⁡𝑛)O(nlogn)입니다.
• 리스트의 길이가 길어질수록 정렬에 걸리는 시간도 비례하여 증가하게 됩니다.

효율성 측면

• heapq를 사용하면 매번 리스트를 정렬하는 대신, 힙 연산을 통해 가장 작은 두 요소를 효율적으로 찾고 제거할 수 있습니다. 힙 연산의 시간 복잡도는 𝑂(log⁡𝑛)O(logn)으로, 정렬을 반복적으로 수행하는 것보다 훨씬 효율적입니다.
• 힙은 내부적으로 이진 트리(binary tree) 구조를 사용하여 요소를 저장하며, 이로 인해 요소 삽입 및 제거 시 항상 최소값을 유지합니다.

heapq를 사용한 예제 코드

import heapq

n = int(input())
nums = list(map(int, input().split()))
cost = 0

# 리스트를 힙으로 변환
heapq.heapify(nums)

while len(nums) > 1:
    # 가장 작은 두 숫자 구하기
    n1 = heapq.heappop(nums)
    n2 = heapq.heappop(nums)

    # 드는 비용 업데이트
    cost += (n1 + n2)
    
    # 두 숫자를 더해서 힙에 다시 넣기
    heapq.heappush(nums, n1 + n2)

print(cost)​

 

3. 느낀 점

자료구조를 잘 활용하는 것이 알고리즘 측면에서 얼마나 큰 이득을 줄 수 있는지 다시 한 번 깨달을 수 있는 문제였습니다. 알고리즘 코딩 테스트를 준비하면서, 다양한 자료구조와 알고리즘을 익히는 것이 중요함을 깨달았습니다. 다양한 문제를 풀면서 특정 상황에서 가장 적합한 자료구조를 선택하는 능력을 키우는 것이 필요합니다. 이 문제를 통해 힙 자료구조의 유용성을 경험하며, 다른 문제에서도 적절히 활용할 수 있는 자신감을 얻었습니다.

 

앞으로도 다양한 문제를 풀고, 효율적인 알고리즘을 찾는 연습을 지속해 나갈 것입니다. 이를 통해 더 나은 문제 해결 능력을 갖추고, 알고리즘 코딩 테스트에서 좋은 결과를 얻을 수 있도록 노력하겠습니다..!