けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder ABC 245 C - Choose Elements (茶色, 300 点)

AtCoder AtCoder300点 茶色diff ABC-C DP 数列 (定数)×Nのグリッド ナップサックDP グラフ 絶対値やminを扱う問題 隣接2要素に制約や目的関数のある問題 【問題集】DPの入門 DP状態:on/off Yes/No判定問題 解空間:O(2^N)通りの選択肢

EDPC C - Vacation と良く似た問題だと思う！！

あと、幅が狭いグリッドでは DP が疑われることが結構多い！

問題概要

長さが N の数列が 2 つ ( A_{0}, \dots, A_{N-1} と B_{0}, \dots, B_{N-1}) 与えられます。

各 i = 0, 1, \dots, N-1 に対して、 A_{i} と B_{i} のいずれかを選ぶことで、新たに数列 W_{0}, \dots, W_{N-1} を作ります。

こうしてできる数列が i = 0, 1, \dots, N-2 に対して |W_{i} - W_{i+1}| \ge K を満たすようにできるかどうかを判定してください。

制約

  • 1 \le N \le 2 \times 10^{5}
  • 0 \le K \le 10^{9}

考えたこと

たとえば A = (9, 8, 3, 7, 2)、 B = (1, 6, 2, 9, 5)、 K = 4 に対しては、下図のようなグラフに対して、左側から出発して右側へと到達できる方法があるかどうかを判定する問題だと言えます。

K = 4 なので、差が 4 以下の部分にのみ矢印を引いています。

Union-Find は嘘解法

「左側と右側がつながれば Yes、途切れたら No」と考えると、一瞬 Union-Find が思い浮かぶかもしれません。

しかし Union-Find を用いると、下図のグラフに対して Yes と返してしまいます (正解は No です)。一旦戻るのは禁止ですが、 Union-Find でそれを禁止するのは困難です。

DP へ

今回のグラフは「左から右へと一方向に流れていくグラフ」です。こういうグラフの経路を扱う問題では DP が有効です。たとえば次のような DP テーブルを用意しましょう。

  • dpA[i]：上側の i 番目の頂点に到達できれば true、そうでなければ false
  • dpB[i]：下側の i 番目の頂点に到達できれば true、そうでなければ false

まず最初は左から 0 番目の頂点は上下ともに到達できる (開始時点で到達している) ので、

dpA[0] = dpB[0] = true;

とします。そして、 i 番目の頂点から i+1 番目の頂点への遷移を考えます。

上側の i 番目の頂点からの遷移

もし dpA[i] = true ならば、

  • 上側の i 番目の頂点から、上側の i+1 番目の頂点へ移動できるとき、dpA[i+1] = true と更新します
  • 上側の i 番目の頂点から、下側の i+1 番目の頂点へ移動できるとき、dpB[i+1] = true と更新します

dpA[i] = false のときは何もしません

下側の i 番目の頂点からの遷移

もし dpB[i] = true ならば、

  • 下側の i 番目の頂点から、上側の i+1 番目の頂点へ移動できるとき、dpA[i+1] = true と更新します
  • 下側の i 番目の頂点から、下側の i+1 番目の頂点へ移動できるとき、dpB[i+1] = true と更新します

dpB[i] = false のときは何もしません

　 そして最後に、dpA[N-1] と dpB[N-1] のいずれかが true ならば "Yes"、そうでなければ "No" です。

　

コード

計算量は O(N) となります。

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

int main() {
    // 入力
    long long N, K;
    cin >> N >> K;
    vector<long long> 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
    vector<bool> dpA(N, false), dpB(N, false);
    dpA[0] = dpB[0] = true;
    for (int w = 0; w < N-1; ++w) {
        // A の w 番目まで到達できるとき、さらに伸ばす
        if (dpA[w]) {
            if (abs(A[w] - A[w+1]) <= K) dpA[w+1] = true;
            if (abs(A[w] - B[w+1]) <= K) dpB[w+1] = true;
        }
        // B の w 番目まで到達できるとき、さらに伸ばす
        if (dpB[w]) {
            if (abs(B[w] - A[w+1]) <= K) dpA[w+1] = true;
            if (abs(B[w] - B[w+1]) <= K) dpB[w+1] = true;
        }
    }

    // A, B のどちらかの N-1 番目まで到達できれば Yes
    if (dpA[N-1] || dpB[N-1]) cout << "Yes" << endl;
    else cout << "No" << endl;
}