[백준] 2616. 소형기관차

문제 링크

[백준] 2616. 소형기관차

풀이 과정

다이나믹 프로그래밍 문제입니다. Top-down 과 Bottom-up 두 방법으로 풀었습니다.

Top-down

sum = new int[N + 1];
dp = new int[N + 1][3];

sum[i] 는 i번 요소까지의 부분합을 의미합니다. dp[i][j] 는 j개의 기관차를 이용해 i번 객차까지지 진행했을 때 태울 수 있는 최대 손님을 의미합니다.

for (int i = 1; i <= N; i++) sum[i] = sum[i - 1] + arr[i - 1];
for (int i = 1; i <= N; i++) Arrays.fill(dp[i], -1);

sum 은 i번 객차까지의 부분합을 의미합니다. 이전 객차까지의 모든 손님에 현재 객차의 손님을 더해 저장해줍니다. dp 배열 요소의 0 값은 손님을 하나도 못태운 것을 의미하는데에 사용되므로, 메모이제이션이 안됐음을 -1로 표현 해줬습니다.

static int go(int idx, int cnt) {
    if (cnt == 3) return 0;
    if (idx >= N) return 0;
    if (dp[idx][cnt] != -1) return dp[idx][cnt];
    if (idx + K - 1 <= N)
        dp[idx][cnt] = Math.max(go(idx + 1, cnt), sum[idx + K - 1] - sum[idx - 1] + go(idx + K, cnt + 1));

    return dp[idx][cnt];
}

go 메소드 는 최대 깊이까지 이동해 태울 수 있는 승객을 계산한 후, dp 배열에 저장하고 그 값을 리턴합니다.

• cnt가 3이면, 3개의 기관차를 다 사용한 것이므로 더 이상 승객을 태울 수 없기에 0을 리턴합니다.
• idx가 N 이상이면, 모든 객차의 순회가 끝났으므로 위와 마찬가지로 더 이상 승객을 태울 수 없게 됩니다.
• dp 요소 값이 -1이 아니면, 그 값을 가져다 사용하기 위해 리턴합니다.
• dp 요소 값이 -1이면, 값을 구한 적이 없기 때문에, 아래 두 값을 비교하여 큰 값을 저장후 리턴합니다.
  • 이번 인덱스를 시작으로 하는 객차들의 승객을 현재 기관차에 태우지 않는 경우와,
  • 이번 인덱스를 시작으로 하는 객차들의 승객을 현재 기관차에 태우는 경우를 비교해 더 많은 승객을 태울 수 있는 경우의 값을 dp[idx][cnt] 에 저장합니다.

Bottom-up

Top-down 방식의 경우, Node.js 로 풀이하면 위와 같이 StackSizeExceeded 라는 에러를 마주하게 됩니다. 이는, 변수의 저장 공간을 널널하게 잡는 JavaScript의 특성상, 호출 스택이 깊어지면 이에 따라 스택이 고갈해버리는 현상이라고 합니다.

📕참고:
https://kscodebase.tistory.com/409

JS로 풀기 위해, Bottom-up 방식으로도 생각해봤습니다.

const sum = Array(N + 1).fill(0);
const dp = Array.from(Array(4), () => Array(N + 1).fill(0));

for (let i = 1; i <= N; i++) sum[i] = sum[i - 1] + arr[i - 1];

누적합을 의미하는 sum 과 메모이제이션을 위한 dp 를 Top-down 과 동일하게 선언해줍니다.

for (let cnt = 1; cnt <= 3; cnt++) {
    for (let idx = cnt * K; idx <= N; idx++) {
        dp[cnt][idx] = Math.max(dp[cnt][idx - 1], dp[cnt - 1][idx - K] + sum[idx] - sum[idx - K]);
    }
}

현재 인덱스를 시작으로 K개의 객차의 승객들을 태우지 않는 경우와, 태우는 경우 중 더 많은 손님을 태울 수 있는 경우의 값을 저장해줍니다.

코드

Top-down

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.Arrays;

public class Main {
    static int N, K;
    static int[] sum;
    static int[][] dp;

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        N = stoi(br.readLine());
        int[] arr = Arrays.stream(br.readLine().split(" ")).mapToInt(Integer::parseInt).toArray();
        K = stoi(br.readLine());
        sum = new int[N + 1];
        dp = new int[N + 1][3];

        for (int i = 1; i <= N; i++) sum[i] = sum[i - 1] + arr[i - 1];
        for (int i = 1; i <= N; i++) Arrays.fill(dp[i], -1);

        System.out.println(go(1, 0));
    }

    static int go(int idx, int cnt) {
        if (cnt == 3) return 0;
        if (idx >= N) return 0;
        if (dp[idx][cnt] != -1) return dp[idx][cnt];
        if (idx + K - 1 <= N)
            dp[idx][cnt] = Math.max(go(idx + 1, cnt), sum[idx + K - 1] - sum[idx - 1] + go(idx + K, cnt + 1));

        return dp[idx][cnt];
    }

    static int stoi(String s) {
        return Integer.parseInt(s);
    }
}

Bottom-up

const fs = require("fs");
const stdin = (
    process.platform === "linux"
        ? fs.readFileSync("/dev/stdin")
        : `7
35 40 50 10 30 45 60
2`
)
    .toString()
    .trim()
    .split("\n");
const input = (() => {
    let line = 0;
    return () => stdin[line++];
})();

const N = +input();
const arr = input().split(" ").map(Number);
const K = +input();
const sum = Array(N + 1).fill(0);
const dp = Array.from(Array(4), () => Array(N + 1).fill(0));

for (let i = 1; i <= N; i++) sum[i] = sum[i - 1] + arr[i - 1];

for (let cnt = 1; cnt <= 3; cnt++) {
    for (let idx = cnt * K; idx <= N; idx++) {
        dp[cnt][idx] = Math.max(dp[cnt][idx - 1], dp[cnt - 1][idx - K] + sum[idx] - sum[idx - K]);
    }
}

console.log(dp[3][N]);