금광 코드 및 해설 (파이썬)

출처: 이것이 취업을 위한 코딩 테스트다 with 파이썬, Q 31, p 375

 

이전에 풀이한 등굣길 문제와 유사하게 다이나믹 프로그래밍(DP)을 사용해 풀이했습니다. 

입력으로 받은 금광에서 1열의 값은 그대로 두고, 2열부터 마지막 열까지 다음을 반복합니다.

첫 행의 경우 현재 위치를 왼쪽, 왼쪽 아래에서부터 접근해올 수 있습니다. 왼쪽 위는 금광이 존재하지 않습니다.

마지막 행의 경우 현재 위치를 왼쪽, 왼쪽 위에서부터 접근해올 수 있습니다. 왼쪽 아래는 금광이 존재하지 않습니다.

중간 행의 경우 현재 위치를 왼쪽, 왼쪽 위, 왼쪽 아래 모두에서 접근해올 수 있습니다. 

접근해올 수 있는 모든 값들 중 최댓값을 현재 금광에서 얻을 수 있는 값과 더해줍니다. 

마지막으로 마지막 열의 모든 값들 중 최댓값을 반환합니다. 

 

import fileinput

# 입력 처리하기
for i, line in enumerate(fileinput.input()):
    if i == 0:
        # 테스트 케이스 개수
        N = int(line)
        tests = []
        sizes = []
    elif i % 2 == 1:
        # 금광 크기
        a, b = list(map(int, line.split()))
        sizes.append((a,b))
        tests.append([[0]*b for _ in range(a)])
    else:
        a, b = len(tests[-1]), len(tests[-1][0])
        vals = list(map(int, line.split()))
        for row, i in enumerate(range(0, len(vals), b)):
            tests[-1][row] = vals[i:i+4]


for test, size in zip(tests, sizes):
    a, b = size
    # 1열, 2열, ... b열까지 
    for j in range(b):
        for i in range(a):
            # 첫 열
            if j == 0:
                continue
            # 첫 행 (왼쪽, 왼쪽 아래에서 접근 가능)
            elif i == 0:
                test[i][j] += max(test[i][j-1], test[i+1][j-1])
            # 마지막 행 (왼쪽 위, 왼쪽에서 접근 가능)
            elif i == (a-1):
                test[i][j] += max(test[i-1][j-1], test[i][j-1])
            # 중간 행 (왼쪽 위, 왼쪽, 왼쪽 아래에서 접근 가능)
            else:
                test[i][j] += max(test[i-1][j-1], test[i][j-1], test[i+1][j-1])

    # 마지막 열의 값들 중 최댓값
    print(max([row[-1] for row in test]))