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

파이썬으로 K가지 색상의 울타리 페인팅 최소 비용 구하기 (동적 계획법 풀이)


문제 정의

N개의 울타리를 K가지 서로 다른 색상으로 칠한다고 가정해 보겠습니다. 단, 인접한 두 울타리는 반드시 다른 색이어야 하며, 전체 비용은 최소화해야 합니다. N×K 크기의 행렬에서 n번째 행, k번째 열의 값은 n번째 울타리를 k번째 색으로 칠할 때 드는 비용을 나타냅니다. 목표는 이 조건을 만족하는 최소 총비용을 구하는 것입니다.

예를 들어 입력이 다음과 같다고 해보겠습니다.

645
327
345
544

이 경우 출력은 14입니다. 첫 번째 울타리부터 차례대로 색상 인덱스 5 → 2 → 3 → 4를 선택하면 인접한 울타리가 같은 색이 되지 않으면서 총비용 14를 달성할 수 있기 때문입니다.

풀이 접근 방식: 동적 계획법

이 문제는 동적 계획법(DP)으로 효율적으로 해결할 수 있습니다. 핵심 아이디어는 각 단계에서 "최소 비용"뿐 아니라 "두 번째로 작은 비용"까지 함께 추적하는 것입니다. 직전 울타리에 이미 사용된 색과 현재 후보 색이 겹치는 경우에는 최솟값을 그대로 사용할 수 없으므로, 두 번째 최솟값으로 대체해야 하기 때문입니다.

구체적인 알고리즘은 다음과 같습니다.

  • n := 행렬의 행(울타리) 개수
  • fc := 0, ft := 0 — 최적(첫 번째) 색상 인덱스와 해당 누적 비용
  • sc := 1, st := 0 — 차선(두 번째) 색상 인덱스와 해당 누적 비용
  • 행렬의 각 행에 대해 다음을 수행합니다.
    • nfc := -1, nft := 무한대 / nsc := -1, nst := 무한대로 초기화
    • 행의 각 인덱스 i와 비용 t에 대해:
      • ct := t + (i ≠ fc이면 ft, 같으면 st)
      • ct ≤ nft이면: nsc := nfc, nst := nft로 기존 값을 밀어낸 뒤 nfc := i, nft := ct로 갱신
      • 그 외에 ct ≤ nst이면: nsc := i, nst := ct로 갱신
    • 행 처리가 끝나면 fc := nfc, ft := nft / sc := nsc, st := nst로 갱신
  • 모든 행을 처리한 후 ft(최소 누적 비용)를 반환합니다.

예제 구현

아래 파이썬 코드로 위 알고리즘을 직접 확인해 볼 수 있습니다.

class Solution:
    def solve(self, matrix):
        n = len(matrix)
        fc, ft = 0, 0
        sc, st = 1, 0
        inf = int(1e18)
        for row in matrix:
            nfc, nft = -1, inf
            nsc, nst = -1, inf
            for i, t in enumerate(row):
                ct = t + (ft if i != fc else st)
                if ct <= nft:
                    nsc, nst = nfc, nft
                    nfc, nft = i, ct
                elif ct <= nst:
                    nsc, nst = i, ct
            fc, ft = nfc, nft
            sc, st = nsc, nst
        return ft

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

입력

[
    [6, 4, 5],
    [3, 2, 7],
    [3, 4, 5],
    [5, 4, 4]
]

출력

14

복잡도 분석

이 알고리즘은 행렬의 모든 원소를 한 번씩만 확인하므로 시간 복잡도는 O(N×K)입니다. 또한 최적/차선 해를 저장하는 상수 개수의 변수만 사용하므로 공간 복잡도는 O(1)입니다. 모든 색상 조합을 일일이 검사하는 완전 탐색과 달리, 이 방식은 선형 시간 안에 최적해를 구할 수 있다는 점이 큰 장점입니다.