なんか久々なせいか計算量解析がガバガバだったのでメモ。

問題概要

正整数$M,K$と長さ$N$の数列$A, B$が与えられます。$(1, 2, \cdots, N)$の連続するとは限らない部分列$P=(P_1, P_2, \cdots, P_{|P|})$であって、以下の条件を満たす物のうち$\sum_{i=1}^{||P|} A_{P_i}$の最大値を求めてください。

• $\sum_{i=1}^{||P|} B_{P_i} \leq M$
• $i=1,2,\cdots, P_{|P|} -1$ について、$P_{i+1} - P_i \leq K$

問題リンク

解法

どう見てもDPなのでDPテーブルを考える。$M$が最大で200と小さいので、そこまでのBの和を持っておくのが自然。$\text{dp}_{i,j} :=$ 条件を満たす列であって最後に$i$を使い、$B$の和が$j$であるような列のうち$A_i$の和としてあり得る最大値、とするDPをすればよい。

空間は$O(NM)$。遷移は、$dp_{i,j}$へと遷移可能なのは$k=1,2,\cdots,K$について$\text{dp}_{i-k,j-B_i}$の$K$通りしか無いので$O(K)$で良い。全体で$O(NMK)$で、余裕を持って通る。

提出(Rust, 7ms)

use library::utils::{chlibs::ChLibs, input::Input, iterlibs::collect::CollectIter};

fn solve(ip: &mut Input) {
    let (n, m) = ip.pair::<usize>();
    let k = ip.next::<usize>();
    let v = (0..n)
        .map(|_| (ip.next::<i64>(), ip.next::<usize>()))
        .collect_vec();

    let mut dp = vec![vec![0i64; m + 1]; n];

    for (i, &(a, b)) in v.iter().enumerate() {
        let mut mn = vec![0; m + 1];

        for j in if i < k { 0 } else { i - k }..i {
            for l in 0..=m {
                mn[l].chmax(dp[j][l]);
            }
        }

        for c in 0..m {
            if c + b <= m {
                dp[i][c + b].chmax(mn[c] + a);
            } else {
                break;
            }
        }
    }

    println!("{}", dp.iter().flatten().max().unwrap());
}

---

自省用にアホのメモをペタリ