[Python/파이썬] 백준 17404번 RGB거리 2

17404번: RGB거리 2

문제

RGB거리에는 집이 N개 있다. 거리는 선분으로 나타낼 수 있고, 1번 집부터 N번 집이 순서대로 있다.

집은 빨강, 초록, 파랑 중 하나의 색으로 칠해야 한다. 각각의 집을 빨강, 초록, 파랑으로 칠하는 비용이 주어졌을 때, 아래 규칙을 만족하면서 모든 집을 칠하는 비용의 최솟값을 구해보자.

• 1번 집의 색은 2번, N번 집의 색과 같지 않아야 한다.
• N번 집의 색은 N-1번, 1번 집의 색과 같지 않아야 한다.
• i(2 ≤ i ≤ N-1)번 집의 색은 i-1, i+1번 집의 색과 같지 않아야 한다.

입력

첫째 줄에 집의 수 N(2 ≤ N ≤ 1,000)이 주어진다. 둘째 줄부터 N개의 줄에는 각 집을 빨강, 초록, 파랑으로 칠하는 비용이 1번 집부터 한 줄에 하나씩 주어진다. 집을 칠하는 비용은 1,000보다 작거나 같은 자연수이다.

출력

첫째 줄에 모든 집을 칠하는 비용의 최솟값을 출력한다.

코드

• DFS (시간 초과)

import sys
input = sys.stdin.readline

n = int(input())
costs = [[] for _ in range(n+1)]
for i in range(n):
    r, g, b = map(int, input().split())
    costs[i+1] = (sorted([(r,0),(g, 1),(b, 2)]))    # 0: red, 1: green, 2: blue

min_cost = 1000001    
color = [-1] * (n + 1)

def dfs(num, cost):
    global min_cost
    
    for c, rgb in costs[num+1]:
        if cost + c > min_cost:
            continue
        if num == 0:    # 1번 집
            color[num+1] = rgb
            dfs(1, c)
            color[num+1] = 0
        elif num == n - 1:  # n번 집
            if rgb == color[1] or rgb == color[num]: 
                continue
            else:
                min_cost = min(min_cost, cost + c)
                return
        else:   # i번 집
            if rgb == color[num]:
                continue
            else:
                color[num+1] = rgb
                dfs(num+1, cost + c)
                color[num+1] = 0

dfs(0, 0)
print(min_cost)

 

DFS와 백트래킹을 생각하고 위와 같이 코드를 작성하였는데 44% 쯤에서 시간초과가 나서 통과할 수 없다. 그래서 다른 방법을 찾아보았다.

우선 위의 코드를 설명하자면

1. n번째 집의 색을 저장하는 배열 color에 n번째 집의 색을 저장하며 dfs를 재귀적으로 돈다.
2. 1번째 집부터 순차적으로 확인하므로 본인 앞 순서의 집과 겹치지 않는지만 확인하면 된다.
3. n번째 집까지 간다면 그 비용을 min_cost와 비교하여 둘 중 최소인 값을 min_cost에 저장한다.
4. 집을 칠하는 비용이 말단 노드까지 가기도 전에 이전에 계산되어 저장된 최소비용인ㅜ min_cost보다 작다면 더이상 볼 필요가 없기 때문에 continue한다.

 

• DP

import sys
input = sys.stdin.readline

n = int(input())
costs = []
min_cost = 1000001

for i in range(n):
    r, g, b = map(int, input().split())
    costs.append([r,g,b])    # 0: red, 1: green, 2: blue

for first in range(3):
    dp = [[1000001 for _ in range(3)] for _ in range(n)]
   
    dp[0][first] = costs[0][first]
    
    for i in range(1, n):
        dp[i][0] = min(dp[i-1][1], dp[i-1][2]) + costs[i][0]
        dp[i][1] = min(dp[i-1][0], dp[i-1][2]) + costs[i][1]
        dp[i][2] = min(dp[i-1][0], dp[i-1][1]) + costs[i][2]

    for j in range(3):
        if j == first:
            continue
        else:
            min_cost = min(min_cost, dp[n-1][j])
            
print(min_cost)

 

1. 첫번째 집이 r, g, b일 때의 케이스를 각각 나누어 구할 것이다.
   1. 1000001로 초기화한 2차원 배열 dp를 만들어준다.
   2. dp[0][first] = costs[0][first]로 첫번째 집의 비용을 저장해준다. 첫번째 집의 색은 고정시켜놓고 생각할 예정이기에 다른 색은 굳이 저장하지 않았다.
   3. 2번째 집부터 n번째 집까지 dp 배열에 소모되는 비용을 저장해나가며 구해준다. 이때 이전 색과 같은 색은 칠할 수 없고 최소 비용을 구하는 것이기에 아래와 같이 코드를 작성한 것이다.
      
      for i in range(1, n):
              dp[i][0] = min(dp[i-1][1], dp[i-1][2]) + costs[i][0]
              dp[i][1] = min(dp[i-1][0], dp[i-1][2]) + costs[i][1]
              dp[i][2] = min(dp[i-1][0], dp[i-1][1]) + costs[i][2]
   4. n번째 집까지 dp 배열의 값을 다 구했다면 n번째 집과 1번째 집의 색이 같지 않은 경우 중 최소값을 min_cost 변수에 저장해둔다.
2. 1번째 집의 색 3가지 경우에 대해 위의 과정을 모두 실행하고 min_cost에 최종적으로 저장된 값이 우리가 구하는 최소 비용이다.