[Flipkart인터뷰] 금광 (python)

소냐. 2024. 2. 11. 10:06

문제

n x m 크기의 금광이 있습니다. 금광은 1x1 크기의 칸으로 나누어져 있으며, 각 칸은 특정한 크기의 금이 들어 있습니다. 채굴자는 첫 번째 열부터 출발하여 금을 캐기 시작합니다. 맨 처음에는 첫번째 열의 어느 행에서든 출발할 수 있습니다. 이후에 m번에 걸쳐서 매번 오른쪽 위, 오른쪽, 오른쪽 아래 3가지 중 하나의 위치로 이동해야 합니다. 결과적으로 채굴자가 얻을 수 있는 금의 최대 크기를 출력하는 프로그램을 작성하세요.

만약 다음과 같이 3 x 4 크기의 금광이 존재한다고 가정합시다.

1	3	3	2
2	1	4	1
0	6	4	7

가장 왼쪽 위의 위치를 (1,1), 가장 오른쪽 아래의 위치를 (n,m)이라고 할때,

위 예시에서는 (2,1) -> (3,2) -> (3,3) -> (3,4)의 위치를 이동하면 총 19만큼의 금을 채굴할 수 있으며, 이때의 값이 최댓값입니다.

입력 조건

첫째 줄에 테스트 케이스 T가 입력됩니다.( 1 <= T <= 1000)

매 테스트 케이스 첫째 줄에 n과 m이 공백으로 구분되어 입력됩니다. (1 <= n, m <= 20)

둘째 줄에 n x m개의 위치에 매장된 금의 개수가 공백으로 구분되어 입력됩니다. (0 <= 각 위치의 매장된 금의 개수 <= 100)

출력 조건

테스트 케이스마다 채굴자가 얻을 수 있는 금의 최대 크기를 출력합니다. 각 테스트 케이스는 줄 바꿈을 이용해 구분합니다.

입력 예시

2
3 4
1 3 3 2 2 1 4 1 0 6 4 7
4 4
1 3 1 5 2 2 4 1 5 0 2 3 0 6 1 2

출력 예시

19
16

풀이

나는 이 문제에서 갈 수 있는 방향이 '오른쪽 위, '오른쪽', '오른쪽 아래'라고 할 때, 방향벡터리스트를 3개 만들어서 각각 움직이면서 3개의 값 중 뭐가 더 큰지 비교해야겠다고 생각했다. 그리고 입력받은 금광 데이터만큼 DP테이블을 만들어서 이동할 때마다 무엇이 최댓값인지 비교해서 최댓값으로 갱신한다.

풀이1)

t = int(input())
for test in range(t):
    n, m = map(int, input().split())
    arr_list = list(map(int, input().split()))
    arr = [[arr_list[i*m+j] for j in range(m)] for i in range(n)]
    dp = [[0 for j in range(m)] for i in range(n)]

    for i in range(n):
        dp[i][0] = arr[i][0]

    dx = [1, 0, -1]
    dy = [1, 1, 1]
    for k in range(m-1):
        for i in range(n):
            x, y = i, k
            for j in range(3):
                nx = x + dx[j]
                ny = y + dy[j]
                if nx < 0 or nx > n-1 or ny < 0 or ny > m-1:
                    continue
                dp[nx][ny] = max(dp[x][y] + arr[nx][ny], dp[nx][ny])

    max_dp = 0
    for i in range(n):
        max_dp = max(max_dp, dp[i][m-1])
    print(max_dp)

그래서 위처럼 코드를 작성했는데, DP를 이용하면서도 반복문을 살펴보면 BFS같기도 하고 그렇다.

내가 작성한 코드의 DP 점화식은 d(n) = max( d(n), d(n-1)+arr(n) ) 인데 d(n)이라는 자기 자신과 max를 비교하는 이유는 움직일 수 있는 방향 3개를 반복문으로 처리해주었기 때문이다. 총 3개의 움직임이 차례로 d(n)에 저장되기 때문에 그 다음 움직임이 전의 움직임보다 더 많은 금광을 캘 수 있는지 비교해야 한다.

위 풀이는 현재 머무는 칸(x, y)에서부터 다음으로 움직일 칸(nx, ny)들의 예상 최댓값을 비교해 갱신해주는 방식이다. 그래서 BFS랑 비슷하다고 느끼는 것 같은데, 이번에는 현재 머무는 칸의 예상 최댓값을 찾기 위해 이전 칸들 중 최댓값을 가져와서 현재 칸의 금광수와 더하는 방식으로 풀어보았다.

이 풀이의 점화식을 세우면 다음과 같다.

d[x, y] = max(d[x-1, y-1], d[x, y-1], d[x+1, y-1]) + arr[x, y]

현재 머무는 칸 x, y 에서

이전 칸들(왼쪽 위, 왼쪽, 왼쪽 아래) 중 최댓값을 찾아 현재 머무는 칸의 금광수와 더한다.

풀이2)

t = int(input())
for test in range(t):
    n, m = map(int, input().split())
    arr_list = list(map(int, input().split()))
    arr = [[arr_list[i*m+j] for j in range(m)] for i in range(n)]
    dp = [[0 for j in range(m)] for i in range(n)]

    for i in range(n):
        dp[i][0] = arr[i][0]

    dx = [1, 0, -1]
    dy = [1, 1, 1]
    for y in range(1, m):
        for x in range(n):
            if x < 1:
                first = 0
            else:
                first = dp[x-1][y-1]
            second = dp[x][y-1]
            if x >= n-1:
                third = 0
            else:
                third = dp[x+1][y-1]
            dp[x][y] = max(first, second, third) + arr[x][y]

    max_dp = 0
    for i in range(n):
        max_dp = max(max_dp, dp[i][m-1])
    print(max_dp)

다만 위 점화식에서 x-1, x+1일 경우 x의 범위에 따라 인덱스 범위를 벗어날 수도 있기 때문에,

범위를 확인해주고 벗어난다면 0으로 설정해준다.

회고

DP문제를 풀 때에는 점화식을 제대로 세우는 것이 중요하다고 한다.

그런데 나는 제대로 점화식을 정하고 푸는 게 아니라 풀다가 정하곤 한다. 이런 점이 나중에는 실수를 만들 것 같기도 하고, 풀이에 확신을 주지 못하는 안 좋은 습관이 아닐까 생각이 든다. 그래서 DP문제인 것 같다는 생각이 들면 점화식을 제대로 세우고 코드를 작성하는 연습을 해야 겠다.

알고리즘 분류

다이나믹 프로그래밍

728x90