[백준] 3114 - 사과와 바나나

https://www.acmicpc.net/problem/3114

3114번: 사과와 바나나

 

문제를 보고 prefix sum + dp이라는 것을 알아차리고 간단하게 해결하려 했지만 WA를 받았다. 먼저 처음 시도한 방법부터 알아보자.

 

map[i][j] = 입력

asum[i][j] = j열의 i행 까지의 사과나무의 누적합(map[1][j] + map[2][j] + ... map[i][j])

bsum[i][j] = j열의 i행 까지의 바나나나무의 누적합(map[1][j] + map[2][j] + ...map[i][j])

dp[i][j] = (i,j)에 도달했을 때 얻을 수 있는 나무의 합의 최댓값

라고 했을 때, dp[i][j]는 3가지의 경우로 채워질 수 있다.

 

• (i - 1, j - 1)에서 (i, j)로 도달하는 경우
  • (i, j) 위의 바나나 나무와 (i, j) 아래의 사과나무를 누적합으로 구하면 된다.
• (i, j - 1)에서 (i, j)로 도달하는 경우
  • (i, j) 위의 바나나 나무와 (i, j) 아래의 사과 나무를 누적합으로 구하면 된다.
• (i - 1, j)에서 (i, j)로 도달하는 경우
  • (i, j)가 사과나무일 경우, dp[i - 1][j]에서 map[i][j]만큼 빼면 된다.

 

이렇게만 구현한 결과가 아래와 같다.

#include <iostream>
using namespace std;

int n, m, map[1501][1501], dp[1501][1501], asum[1501][1501], bsum[1501][1501];

int main() {
    ios_base::sync_with_stdio(0);
    cin.tie(0);

    cin >> n >> m;
    for (int i = 1; i <= n; i++) {
        for (int j = 1; j <= m; j++) {
            char a;
            int b;
            cin >> a >> b;
            asum[i][j] = asum[i - 1][j];
            bsum[i][j] = bsum[i - 1][j];
            if (a == 'A') {
                map[i][j] = b;
                asum[i][j] += b;
            } else {
                map[i][j] = -b;
                bsum[i][j] += b;
            }
        }
    }
    for (int i = 1; i <= n; i++) {
        for (int j = 1; j <= m; j++) {
            dp[i][j] = dp[i - 1][j - 1] + (asum[n][j] - asum[i][j]) + (bsum[i - 1][j]);
            dp[i][j] = max(dp[i][j], dp[i][j - 1] + (asum[n][j] - asum[i][j]) + (bsum[i - 1][j]));
            if (map[i][j] > 0) {
                dp[i][j] = max(dp[i][j], dp[i - 1][j] - map[i][j]);
            } else {
                dp[i][j] = max(dp[i][j], dp[i - 1][j]);
            }
        }
    }
    cout << dp[n][m];
}

 

하지만 틀린 답이었다. 어디가 틀렸을까? 위 코드로 문제 예제 1번을 돌리면 dp 배열은 다음과 같다.

5	9	13
4	12	20
2	8	20
2	8	17

 

이것이 과연 맞는 것일까?

입력이 다음과 같다고 생각해 보자.

4 1
B2
A3
A2
B1

답은 분명히 0인데, dp 배열은 위에서 아래로 5 4 2 2를 나타낼 것이다.

 

결론은, 맨 처음 열에 대해서는 갈 수 있는 과정이 단 하나(아래로만 내려가는 것)이기 때문에 바나나 나무는 얻을 수 없다. 따라서 첫째 열만 미리 전처리 해주고, 두 번째 열부터 dp 점화식을 적용하면 된다.

 

#include <iostream>
using namespace std;

int n, m, map[1501][1501]; // map[i][j] > 0이면 사과, map[i][j] < 0이면 바나나
int dp[1501][1501];
int asum[1501][1501], bsum[1501][1501];

int main() {
    ios_base::sync_with_stdio(0);
    cin.tie(0);

    cin >> n >> m;
    for (int i = 1; i <= n; i++) {
        for (int j = 1; j <= m; j++) {
            char a;
            int b;
            cin >> a >> b;
            asum[i][j] = asum[i - 1][j];
            bsum[i][j] = bsum[i - 1][j];
            if (a == 'A') {
                map[i][j] = b;
                asum[i][j] += b;
            } else {
                map[i][j] = -b;
                bsum[i][j] += b;
            }
        }
    }
    
    // 전처리
    for (int i = 1; i <= n; i++) {
        dp[i][1] = asum[n][1] - asum[i][1];
    }

    for (int i = 1; i <= n; i++) {
        for (int j = 2; j <= m; j++) {
            // 왼쪽 위에서 오는 경우
            dp[i][j] = dp[i - 1][j - 1] + (asum[n][j] - asum[i][j]) + (bsum[i - 1][j]);

            // 왼쪽에서 오는 경우
            dp[i][j] = max(dp[i][j], dp[i][j - 1] + (asum[n][j] - asum[i][j]) + (bsum[i - 1][j]));

            // 위에서 오는 경우, 사과나무의 누적합만 변하게 된다
            if (map[i][j] > 0) {
                // 현재 도달한 칸이 사과나무 일때는 위의 dp값에서 빼줘야 한다
                dp[i][j] = max(dp[i][j], dp[i - 1][j] - map[i][j]);
            } else {
                dp[i][j] = max(dp[i][j], dp[i - 1][j]);
            }
        }
    }
    cout << dp[n][m];
}