파이썬(Python)으로 돌 무더기 병합 최소 비용 구하기 – 구간 DP 풀이

문제 설명

N개의 돌 무더기가 한 줄로 나열되어 있고, i번째 무더기에는 stones[i]개의 돌이 들어 있다고 가정해 보겠습니다. 한 번의 이동은 연속된 K개의 무더기를 하나로 합치는 작업이며, 이때 드는 비용은 해당 K개 무더기에 있는 돌의 총 개수와 같습니다. 목표는 모든 돌 무더기를 하나로 합칠 때 드는 최소 비용을 구하는 것이고, 합치는 것이 불가능한 경우에는 -1을 반환해야 합니다.

예시

입력이 nums = [3, 2, 4, 1], K = 2라고 해봅시다. 이때 출력은 20이 됩니다.

• 처음 상태: [3, 2, 4, 1]
• [3, 2]를 비용 5로 합침 → [5, 4, 1]
• [4, 1]을 비용 5로 합침 → [5, 5]
• [5, 5]를 비용 10으로 합침 → [10]

총 비용은 5 + 5 + 10 = 20이며, 이것이 가능한 최솟값입니다.

접근 방법: 구간 동적 계획법(Interval DP)

이 문제는 구간 단위 동적 계획법으로 효율적으로 해결할 수 있습니다. dp[i][j]를 'i번째부터 j번째 무더기를 최소 비용으로 합치는 비용'이라고 정의한 뒤, 작은 구간의 답을 이용해 점차 넓은 구간의 답을 채워 나갑니다.

여기서 중요한 두 가지 조건이 있습니다.

• 실행 가능 여부 확인: K개의 무더기를 합칠 때마다 전체 무더기 수는 K-1씩 줄어듭니다. 따라서 n개의 무더기를 1개로 만들려면 (n-1)이 (K-1)로 나누어떨어져야만 하며, 그렇지 않으면 -1을 반환합니다.
• 병합 시점의 비용 추가: 구간 [i, j]의 길이가 (K-1)로 나누어떨어질 때, 즉 해당 구간이 하나의 무더기로 완전히 합쳐질 수 있을 때만 그 구간의 돌 전체 합(sums[j+1] - sums[i])을 비용에 더합니다.

알고리즘 단계

• n := nums의 크기
• (n-1) mod (K-1) ≠ 0이면 -1 반환
• dp := n × n 크기의 2차원 배열(0으로 초기화)
• sums := 크기 (n+1)의 누적 합 배열(0으로 초기화)
• i를 1부터 n까지 반복하며 sums[i] := sums[i-1] + nums[i-1] 계산
• 구간 길이 length를 K부터 n까지 반복:
  • 시작점 i를 0부터 n-length까지 반복
  • 끝점 j := i + length - 1
  • dp[i][j] := 무한대로 초기화
  • 분할점 t를 i부터 j-1까지 (K-1)씩 증가시키며 dp[i][j] = min(dp[i][j], dp[i][t] + dp[t+1][j]) 갱신
  • (j-i) mod (K-1) = 0이면 dp[i][j] += sums[j+1] - sums[i]
• 최종 결과 dp[0][n-1] 반환

구현 코드

아래 파이썬 구현을 통해 위 알고리즘이 실제로 어떻게 동작하는지 확인할 수 있습니다.

import heapq

def solve(nums, K):
    n = len(nums)
    if (n-1)%(K-1) != 0:
        return -1
    dp = [[0]*n for _ in range(n)]
    sums = [0]*(n+1)
    for i in range(1, n+1):
        sums[i] = sums[i-1] + nums[i-1]
    for length in range(K, n+1):
        for i in range(n-length+1):
            j = i + length - 1
            dp[i][j] = float('inf')
            for t in range(i, j, K-1):
                dp[i][j] = min(dp[i][j], dp[i][t] + dp[t+1][j])
            if (j-i)%(K-1) == 0:
                dp[i][j] += sums[j+1] - sums[i]
    return dp[0][n-1]

nums = [3, 2, 4, 1]
K = 2
print(solve(nums, K))

입력

[3, 2, 4, 1], 2

출력

20

핵심 포인트 정리

• 각 병합 단계에서 무더기 수가 K-1씩 감소하므로, (n-1) % (K-1) == 0일 때만 모든 돌을 하나로 합칠 수 있습니다.
• 누적 합 배열(sums)을 사용하면 임의 구간의 돌 개수를 O(1)에 구할 수 있습니다.
• 분할점 t를 K-1 간격으로 건너뛰며 탐색해야 두 부분 구간이 각각 유효한 병합 단위가 됩니다.
• 시간 복잡도는 대략 O(n³ / K) 수준으로, 구간 DP의 전형적인 형태를 따릅니다.