Computer >> 컴퓨터 >  >> 프로그래밍 >> Python

파이썬으로 리스트를 하나의 정수로 줄이는 최소 비용 구하기

문제 설명

숫자로 이루어진 리스트 nums가 주어졌다고 가정해 봅시다. 우리는 리스트에서 임의의 두 수를 골라 제거한 뒤, 그 두 수의 합을 리스트 끝에 추가하는 방식으로 리스트의 길이를 줄일 수 있습니다. 이때 각 연산의 비용은 제거한 두 정수의 합입니다. 목표는 nums를 하나의 정수로 줄일 때 드는 최소 총 비용을 구하는 것입니다.

예시로 이해하기

예를 들어 입력이 nums = [2, 3, 4, 5, 6]이라면 결과는 45입니다.

  • 2와 3을 꺼내 합친다 → [4, 5, 6, 5], 비용 5
  • 4와 5를 꺼내 합친다 → [6, 5, 9], 비용 9
  • 6과 5를 꺼내 합친다 → [9, 11], 비용 11
  • 9와 11을 꺼내 합친다 → [19], 비용 20

각 단계의 비용을 모두 더하면 5 + 9 + 11 + 20 = 45가 됩니다.

접근 방법: 그리디 + 최소 힙

이 문제는 허프만 코딩(Huffman Coding)과 유사한 아이디어로 해결할 수 있습니다. 핵심은 작은 수를 자주, 큰 수를 적게 더하는 것입니다. 어떤 수가 합산 과정에 여러 번 참여할수록 총 비용에 여러 번 기여하게 되므로, 작은 값부터 먼저 합치는 것이 유리합니다.

이를 위해 최소 힙(min-heap)을 활용하며, 절차는 다음과 같습니다.

  1. nums의 원소들로 최소 힙을 만든다.
  2. ans = 0으로 초기화한다.
  3. 힙의 크기가 2 이상인 동안 다음을 반복한다.
    • 힙에서 가장 작은 원소 a를 꺼낸다.
    • 힙에서 그다음으로 작은 원소 b를 꺼낸다.
    • ans += a + b
    • a + b를 다시 힙에 삽입한다.
  4. ans를 반환한다.

구현 예제

class Solution:
    def solve(self, nums):
        import heapq
        heapq.heapify(nums)
        ans = 0
        while len(nums) >= 2:
            a = heapq.heappop(nums)
            b = heapq.heappop(nums)
            ans += a + b
            heapq.heappush(nums, a + b)
        return ans

ob = Solution()
nums = [2, 3, 4, 5, 6]
print(ob.solve(nums))

입력

[2, 3, 4, 5, 6]

출력

45

복잡도 분석

힙의 삽입·삭제 연산은 한 번에 O(log n)이 소요되며, 전체 과정에서 n−1번의 병합이 발생하므로 시간 복잡도는 O(n log n)입니다. 공간 복잡도는 힙 저장을 위해 O(n)입니다.