[BOJ/Kotlin] 2240번 자두나무 : DP

문제 설명

자두는 자두를 좋아한다. 그래서 집에 자두나무를 심어두고, 여기서 열리는 자두를 먹고는 한다. 하지만 자두는 키가 작아서 자두를 따먹지는 못하고, 자두가 떨어질 때까지 기다린 다음에 떨어지는 자두를 받아서 먹고는 한다. 자두를 잡을 때에는 자두가 허공에 있을 때 잡아야 하는데, 이는 자두가 말랑말랑하여 바닥에 떨어지면 못 먹을 정도로 뭉개지기 때문이다.

매 초마다, 두 개의 나무 중 하나의 나무에서 열매가 떨어지게 된다. 만약 열매가 떨어지는 순간, 자두가 그 나무의 아래에 서 있으면 자두는 그 열매를 받아먹을 수 있다. 두 개의 나무는 그다지 멀리 떨어져 있지 않기 때문에, 자두는 하나의 나무 아래에 서 있다가 다른 나무 아래로 빠르게(1초보다 훨씬 짧은 시간에) 움직일 수 있다. 하지만 자두는 체력이 그다지 좋지 못해서 많이 움직일 수는 없다.

자두는 T(1≤T≤1,000)초 동안 떨어지게 된다. 자두는 최대 W(1≤W≤30)번만 움직이고 싶어 한다. 매 초마다 어느 나무에서 자두가 떨어질지에 대한 정보가 주어졌을 때, 자두가 받을 수 있는 자두의 개수를 구해내는 프로그램을 작성하시오. 자두는 1번 자두나무 아래에 위치해 있다고 한다.

입력

첫째 줄에 두 정수 T, W가 주어진다. 다음 T개의 줄에는 각 순간에 자두가 떨어지는 나무의 번호가 1 또는 2로 주어진다.

출력

첫째 줄에 자두가 받을 수 있는 자두의 최대 개수를 출력한다.

문제 풀이 아이디어

문제를 간단히 요약하면 T초 동안 최대 W번 만큼 움직일 수 있을 때, 얻을 수 있는 자두의 최대값을 구하면 된다.

이 때 고려해야하는 요소는 아래 3가지이다.
현재 자두(사람)의 위치, 시간, 움직인 횟수

이를 바탕으로 점화식을 세우기 위한 배열을 만들 수 있다.
d[pos][T][W]  

자두나무에서 자두가 떨어질 때 발생할 수 있는 상황은 크게 2가지이다.

1. 1번 자두나무에서 자두가 떨어지는 경우

1-1) 사람 자두가 1번 나무 아래에 있는 경우
d[1][T][W] = max(d[1][T-1][W], d[1][T-1][W-1]) + 1 //자두 획득
1-2) 사람 자두가 2번 나무 아래에 있는 경우
d[2][T][W] = max(d[1][T-1][W-1],d[2][T-1][W]) //자두 획득하지 못하므로 +1 하지 않음.

2. 2번 자두나무에서 자두가 떨어지는 경우

2-1) 사람 자두가 1번 나무 아래에 있는 경우
d[1][T][W] = max(d[1][T-1][W],d[2][T-1][W-1]) //자두 획득하지 못하므로 +1 하지 않음.
2-2) 사람 자두가 2번 나무 아래에 있는 경우
d[2][T][W] = max(d[1][T-1][W-1],d[2][T-1][W]) +1 //자두 획득

점화식을 위같이 세웠을 때 첫 나무 값에 따라 초기값을 세워준다.

1. 자두가 처음 떨어지는 나무가 1번 나무일 때
d[1][1][0] = 1 //자두가 1번 나무 아래에 있고, 1초가 지났을 때, 움직임이 0회일 때 자두 1개를 획득 할 수 있다.
d[2][1][0] = 0 //자두가 2번 나무 아래에 있고, 1초가 지났을 때, 움직임이 0회일 때 자두 0개를 획득 할 수 있다.

2. 자두가 처음 떨어지는 나무가 2번 나무일 때
d[1][1][0] = 0  //자두가 1번 나무 아래에 있고, 1초가 지났을 때, 움직임이 0회일 때 자두 0개를 획득 할 수 있다.
d[2][1][0] = 1 //자두가 2번 나무 아래에 있고, 1초가 지났을 때, 움직임이 0회일 때 자두 1개를 획득 할 수 있다.

그리고 T초를 돌면서 배열의 값을 채워준다. 
현재 자두가 떨어지는 트리의 값을 확인하고 1,2번 상황에 따라 d값을 채워주면 된다.

우리는 t초동안 최대 w만큼 움직여서 얻을 수 있는 최대 자두수가 궁금하기 때문에,
t초동안 (w=) 0번, 1번, 2번 움직여서 자두를 얻는 경우를 모두 확인해 d[lastPos][T][W] 의 max를 구해 답을 도출하면 된다.

코드

import java.io.BufferedReader
import java.io.InputStreamReader
import java.lang.Integer.max
import java.math.*

fun main() = with(BufferedReader(InputStreamReader(System.`in`))) {
    val input = readLine().split(" ")
    val T = input[0].toInt()
    val W = input[1].toInt()

    val tree = Array(T+1) { 0 } //자두가 떨어지는 트리 위치 array
    var d = Array(3) { Array(T + 1) { Array(W + 1) { 0 } } }

    for (i in 1..T) {
        tree[i] = readLine().toInt()
    }

    if (tree[1] == 1) { //첫나무 1
        d[1][1][0] = 1
        d[2][1][1] = 0
    } else { //첫나무 2
        d[1][1][0] = 0
        d[2][1][1] = 1
    } //초기값 세팅

    for (i in 2..T) {
        if (tree[i] == 1) { //현재 자두나무 1
            d[1][i][0] = d[1][i-1][0] + 1 //한번도 안 움직인 경우
            d[2][i][0] = d[2][i-1][0]
            for (j in 1..W) { //w는 이동한 횟수
                d[1][i][j] = max(d[1][i-1][j],d[2][i-1][j-1])+1
                d[2][i][j] = max(d[1][i-1][j-1],d[2][i-1][j])
            }
        } else { // 현재 자두나무 2
            d[1][i][0] = d[1][i-1][0] //한번도 안 움직인 경우
            d[2][i][0] = d[2][i-1][0] + 1
            for (j in 1..W) { //w는 이동한 횟수
                d[1][i][j] = max(d[1][i-1][j],d[2][i-1][j-1])
                d[2][i][j] = max(d[1][i-1][j-1],d[2][i-1][j]) +1
            }
        }
    } //점화식 배열안에 값 다 넣음

    var ans = -1
    for (i in 0 .. W){
        ans = max(d[tree[T]][T][i],ans)
    }

    println(ans)
}

+) 더 나은 코드

    for (i in 2..T) {
        if (tree[i] == 1) { //현재 자두나무 1
            d[1][i][0] = d[1][i-1][0] + 1 //한번도 안 움직인 경우
            for (j in 1..W) { //w는 이동한 횟수
                d[1][i][j] = max(d[1][i-1][j],d[2][i-1][j-1])+1
                d[2][i][j] = max(d[1][i-1][j-1],d[2][i-1][j])
            }
        } else { // 현재 자두나무 2
            for (j in 1..W) { //w는 이동한 횟수
                d[1][i][j] = max(d[1][i-1][j],d[2][i-1][j-1])
                d[2][i][j] = max(d[1][i-1][j-1],d[2][i-1][j]) +1
            }
        }
    } //점화식 배열안에 값 다 넣음

수정된 코드에서는 자두나무가 2 일때 한번도 안 움직인 경우를 계산하지 않았다.
왜냐하면 나무2에서 떨어지는 나무를 줍기위해서는 최소 1번은 이동해야 주울수 있기 때문이다.

처음 자두(사람)의 위치는 1에서 시작하므로 나무2에서 떨어지는 자두를 줍기위해서는 움직여야한다.
즉 움직임이 0인 상태에서 나무2에서 자두를 줍는 경우는 없다.
그러므로 해당 코드는 삭제해도 무방하다.

배운점

dp는 계산한 이전 값들을 모두 저장해 필요한 값을 도출해내는 것이다.
따라서 이 문제에서도 1초부터 T초까지 나타날 수 있는 경우의 수를 모두 저장한 뒤
해당 문제에서 필요한 자두의 max를 값을 저장한 배열에서 찾아내는게 중요하다.

이런 dp의 본질을 잊지말고 문제에 접근해야겠다.