문제 링크
(2024/9/13 기준) Gold V

문제 읽기

문제가 살짝 난해합니다;
$1$초부터 $T$초까지 매 초 $1$번 나무 또는 $2$번 나무에서 하나씩, 총 $T$개의 자두가 떨어집니다. $W$번 이하로 이동하며 떨어지는 자두를 받았을 때의 최댓값을 구해봅시다. 처음 시작은 $1$번 나무입니다.

풀이

제한을 잘 살펴봅시다. $T \leq 1 000$, $W \leq 30$, 나무가 $2$그루 있으니, dp[현재 시간][현재 위치][현재까지 이동한 횟수]로 3차원 DP를 해볼 수 있을 것 같습니다.

우선, DP 테이블에 초깃값들을 채워봅시다. $t$초에 자두가 떨어진 나무의 위치를 $P_t$라고 하면, $1$초에 $1$번 나무에서 시작하니, $DP_{1, 1, 0}=[P_i = 1]$이 되겠고 (편의상 $[c]$는 조건 $c$가 참이면 $1$, 거짓이면 $0$ 값을 가진다고 약속합시다), $1$초에 자두가 떨어지기 전 $2$번 나무로 움직였다면, $DP_{1, 2, 1}=[P_i = 2]$가 됩니다. (한 번 움직였으니, $DP_{1, 2, 0}$이 아닌, $DP_{1, 2, 1}$입니다!)

이후 점화식은 다음과 같이 정의할 수 있습니다.

  • $(k = 0)\ DP_{i, j, 0} = DP_{i-1, j, 0} + [P_i = j]$
  • $(k \neq 0)\ DP_{i, j, k} = \textrm{max}(DP_{i-1, j, k}, DP_{i-1, 3-j, k-1}) + [P_i = j]$

여기서 $j$는 나무의 번호이므로, $1$ 또는 $2$만 들어갈 수 있습니다. 따라서, $3-j$는 반대쪽 나무를 뜻합니다. 점화식이 조금 복잡하지만, $i$초에 $j$번째 나무에 있고, $k$번 이동했을 때 받을 수 있는 자두의 최대 개수를 나타내므로, $i-1$초에서 $j$번 나무에 있는 경우와 나무를 바꾼 경우(즉, $3-j$번 나무에 있는 경우) 중에서 최댓값을 선택하여 갱신합니다. 시간복잡도는 $O(TW)$가 됩니다.

코드

use std::{fmt::Write, io::{stdin, Read}, str::SplitWhitespace};

fn main() {
    let mut buf = String::new();
    let mut out = String::new();

    stdin().read_to_string(&mut buf).unwrap();
    let mut scan = Scanner::new(buf.split_whitespace());
    //let tc = scan.next();
    //for i in 0..tc {
        solve(&mut scan, &mut out, 0);
    //}
    println!("{out}");
}

// overflow, edge case, and 1 based index?
fn solve<'a>(scan: &mut Scanner<'a, SplitWhitespace<'a>>, out: &mut String, _tc: usize) {
    let n: usize = scan.next();
    let w: usize = scan.next();
    let mut a = vec![0; n + 1];
    for i in 1..=n {
        a[i] = scan.next::<i32>();
    }

    let mut dp = vec![vec![vec![-1; w + 1]; 3]; n + 1];
    dp[1][1][0] = if a[1] == 1 { 1 } else { 0 }; dp[1][2][1] = if a[1] == 2 { 1 } else { 0 };
    for i in 2..=n {
        for j in 1..=2 {
            dp[i][j][0] = dp[i - 1][j][0] + if a[i] == j as i32 { 1 } else { 0 };
            for k in 1..=w {
                dp[i][j][k] = i32::max(dp[i - 1][3 - j][k - 1], dp[i - 1][j][k]) + if a[i] == j as i32 { 1 } else { 0 };
            }
        }
    }
    let mut ans = dp[n][1][0];
    for i in 1..=2 {
        for j in 0..=w {
            if ans < dp[n][i][j] { ans = dp[n][i][j]; }
        }
    }
    writeln!(out, "{ans}").ok();
}

struct Scanner<'a, I: Iterator<Item = &'a str>> {
    iter: I,
}

impl<'a, I: Iterator<Item = &'a str>> Scanner<'a, I> {
    fn new(from: I) -> Self {
        Self { iter: from }
    }
    fn next<T: std::str::FromStr>(&mut self) -> T {
        self.iter.next().unwrap().parse().ok().unwrap()
    }
}

평범한 바텀업 3차원 DP 코드입니다. 마지막에 몇 번을 움직여야 최대가 나올 지 알 수 없으니, $DP_{T, 1, 1}$부터 $DP_{T, 2, W}$까지 돌면서 최댓값을 갱신해주면 됩니다.

후기

DP 어렵습니다… 엄청 많이 고민하고 풀었던 문제입니다.