[python] 백준 17404 : RGB 거리 2

문제 링크
https://www.acmicpc.net/problem/17404

접근

먼저 완전탐색을 생각해보았는데, 2^1000 번 이상 연산해야하므로 시간초과가 날 것을 예상했다. 그래서 dp를 이용하여 풀어야겠다고 생각했다.

처음엔 하나의 dp 배열을 이용하여 문제를 해결하려고 했는데, 첫 집과 마지막 집의 색이 달라야한다는 조건때문에 자꾸 원하는 결과를 얻지 못하는 문제점이 생겼다.

그리하여 0으로 시작하는 경우, 1로 시작하는 경우, 2로 시작하는 경우 세가지로 나누어 3개의 dp 배열을 만들어서 풀었다.
dp[root][i][x]
root 는 첫 집을 어떤 색으로 칠하는지(r==0, g==1, b==2) 나눈 것이고,
i번째에 색깔 x(0,1,2)를 칠하는 경우의 최소 합을 저장한 리스트이다.
점화식은 다음과 같다.

dp[root][i][x] = min(dp[root][i - 1][(x + 1) % 3],
		     dp[root][i - 1][(x + 2) % 3]) + rgb[i][x]

정답 코드

import sys

sys.setrecursionlimit(10 ** 8)
input = lambda: sys.stdin.readline().rstrip()

n = int(input())
rgb = [0 for _ in range(n)]

for i in range(n):
    r, g, b = map(int, input().split())
    rgb[i] = [r, g, b]

dp = [[[999999, 999999, 999999] for _ in range(n)] for _ in range(3)]


for root in range(3):
    dp[root][0][root] = rgb[0][root]
    for i in range(1, n):
        for x in range(3):
            if i == n - 1 and root == x:
                continue
            dp[root][i][x] = (
                min(dp[root][i - 1][(x + 1) % 3], dp[root][i - 1][(x + 2) % 3])
                + rgb[i][x]
            )

m = 99999999
for root in range(3):
    for x in range(3):
        m = min(m, dp[root][n - 1][x])
print(m)