Softeer 복잡한 조립라인1 (난이도 4)

정답 코드

import sys
Input = sys.stdin.readline
K, N = map(int, Input().split())

lines = []
W = [[[0 for i in range(K)] for j in range(K)] for o in range(N)]

for i in range(N - 1):
    lines.append(list(map(int, Input().split())))
    for j in range(K):
        cnt = 0
        for o in range(K):
            if o != j:
                W[i][j][o] = lines[i][cnt + K + j * (K - 1)]
                # 선언의 반대로 사용하는 것 헷갈리지 말기
                cnt += 1
lines.append(list(map(int, Input().split())))

dp = [[0 for i in range(K)] for j in range(N)]
for i in range(K):
    dp[0][i] = lines[0][i]

temp = [0 for o in range(K)]
for i in range(1, N):
    for j in range(K):
        for o in range(K):
            if o != j:
                temp[o] = dp[i - 1][o] + W[i - 1][o][j]
            else:
                temp[o] = dp[i - 1][o]
        dp[i][j] = min(temp) + lines[i][j]


print(min(dp[N - 1]))

알고리즘 유형

dynamic programming

후기

문제 설명이 모호해서 어려웠던 문제.
2번째 줄부터 입력 값이 K * K개인 것을 몰라서 오래 걸렸다.