백준 17404번 : RGB거리 2 in Python

17404번: RGB거리 2

 

www.acmicpc.net

이 문제는 원형으로 배열을 이루는 집들을 3가지 색으로 색칠하되, 서로 이웃한 집의 색이 서로 달라야 한다는 조건이 있는 경우 집을 모두 색칠하는데 드는 비용의 최솟값을 구하는 문제이다.

단순히 모든 집을 색칠하는 경우를 고려하면 집의 개수가 N일 때 O(3^N)의 시간이 소요되므로 매우 오래 걸린다. 그러므로 dynamic programming을 이용해 비용의 최솟값을 계산한다.

만약 집이 원형이 아닌 직선 배열을 이루고 있는 경우라면 문제가 간단하다. 우선 첫 번째 집의 색을 하나 정하고 이후의 집들을 순서대로 색칠하면서 앞의 집과 색이 다르도록 색칠해주면 된다. 하지만 이 문제는 집들이 원형으로 배열을 이루고 있기 때문에 마지막 집이 첫 번째 집과 이웃해 두 집의 색이 같으면 안 된다.

그러므로 위 직선 배열의 경우와 비슷하게 dynamic programming을 구현하나 마지막 집까지 구한 비용에서 첫 번째 집과 색이 같은 경우는 제외하면 된다.

이를 코드로 구현하면 다음과 같다.

import sys
input = sys.stdin.readline
N = int(input())
costs = []
for _ in range(N):
    costs.append(list(map(int, input().split())))
# 비용의 최솟값을 저장하는 변수
res = sys.maxsize
# 첫 번째 집의 색에 따라 dynamic programming을 구현한다.
for initColor in range(3):
    dp = [[sys.maxsize]*3 for _ in range(N)]
    # 첫 번째 집의 색에 해당하는 entry만 초기값을 설정한다.
    for j in range(3):
        if initColor == j: dp[0][j] = costs[0][j]
    for i in range(1, N):
        for j in range(3):
            # 현재 집과 다른 색의 이전 집들의 비용의 합의 최솟값 중 더 작은 값을 골라 비용의 최솟값을 구한다.
            dp[i][j] = min(dp[i-1][(j+1)%3], dp[i-1][(j+2)%3]) + costs[i][j]
    # 마지막 집의 색이 첫 번째 집의 색과 다른 경우에 대해서만 최솟값을 구한다.
    for j in range(3):
        if initColor != j: res = min(res, dp[N-1][j])
print(res)