AtCoder dp例題といてみた3

今回の問題

問題文

長さ $N$ の整数からなる数列 $A=(A_1, \ldots, A_N), B=(B_1, \ldots, B_N)$ が与えられます。

以下の条件を全て満たす長さ $N$ の数列 $X=(X_1, \ldots, X_N)$ が存在するかを判定してください。

• すべての $i(1 \leq i \leq N)$ について、$X_i = A_i$ または $X_i = B_i$
• すべての $i(1 \leq i \leq N-1)$ について、$|X_i - X_{i+1}| \leq K$

制約

• $1 \leq N \leq 2 \times 10^5$
• $0 \leq K \leq 10^9$
• $1 \leq A_i, B_i \leq 10^9$
• 入力は全て整数である

回答

#include <bits/stdc++.h>
using namespace std;

int main() {
    int N, K;
    cin >> N >> K;

    vector<int> A(N), B(N);
    for(int i = 0; i < N; i++) cin >> A[i];
    for(int i = 0; i < N; i++) cin >> B[i];

    // dp[i] := i番目で A[i] (または B[i]) を選ぶことが可能か
    // メモリ節約のため、配列全体を持たず「前回の状態」だけ持ちます
    bool can_use_A = true; // 最初はどちらも選べる
    bool can_use_B = true;

    for(int i = 0; i < N - 1; i++) {
        bool next_can_use_A = false;
        bool next_can_use_B = false;

        // 次の A[i+1] を選べるかチェック
        // 前の A[i] から遷移できるか？
        if (can_use_A && abs(A[i] - A[i+1]) <= K) next_can_use_A = true;
        // 前の B[i] から遷移できるか？
        if (can_use_B && abs(B[i] - A[i+1]) <= K) next_can_use_A = true;

        // 次の B[i+1] を選べるかチェック
        if (can_use_A && abs(A[i] - B[i+1]) <= K) next_can_use_B = true;
        if (can_use_B && abs(B[i] - B[i+1]) <= K) next_can_use_B = true;

        // 更新
        can_use_A = next_can_use_A;
        can_use_B = next_can_use_B;
    }

    if (can_use_A || can_use_B) {
        cout << "Yes" << endl;
    } else {
        cout << "No" << endl;
    }

    return 0;
}

自分なりの解説

dp[i][0]: $i$ 番目で $A_i$ を選ぶことが可能か
dp[i][1]: $i$ 番目で $B_i$ を選ぶことが可能か
というふうにすることにより
前回 $A_i$ を選べていて、かつ $|A_i - A_{i+1}| \le K$ である
前回 $B_i$ を選べていて、かつ $|B_i - A_{i+1}| \le K$ である
という漸化式が立てられる。