NoiminのNoise

競技プログラミング (多め) とWeb (たまに) ,自然言語処理 (ブログではまだ)。数式の書き方を一気に KaTeX に変えようとして記事を全削除してインポートし直すなどしたので,過去にブックマークされた記事は URL が変わってしまっている可能性があります…….

Codeforces Round #605 Div. 3 E. Nearest Opposite Parity

問題: https://codeforces.com/contest/1272/problem/E

公式解説: https://codeforces.com/blog/entry/72132

問題概要

長さ $n (\leq 2 \times 10^5) $ の数列 $a_1, a_2, \cdots, a_n (1 \leq a_i \leq n) $ がある.

この数列を使って,ある要素から別の要素への移動を繰り返す.現在 $i $ 番目の要素にいるとすると,移動できる要素は $i-a_i $ 番目の要素または $i+a_i $ 番目の要素である.ただし,数列の範囲をはみ出して移動することはできない.

すべての要素について,以下を求めよ.

  • 着目する要素を $i $ 番目 としたとき,$a_i $ と偶奇が異なる要素 $a_j $ までにたどり着くまでの最小の移動回数

解法概要

$a_i \% 2 = 1 $ となるような $i $ 番目の要素をスタートとして $a_j \% 2 = 0 $ となるような要素への移動を考える ($a_i \% 2 = 0 $ となる要素をスタートとする場合は偶奇を逆にすれば良い) .

この場合, $a_j \% 2 = 0 $ となるような要素であればいずれの要素もゴールになりうる.また,答えである移動回数を記録したいのは $j $ に対してではなくて $i $ に対してである.そのため,逆辺を張ったグラフを作って BFS を行う.ただし,普段の BFS は始点が 1 つであるが今回は $a_j \% 2 = 0 $ となるような要素をすべて始点としてキューに突っ込む.そうすることで,スタートとゴールの対応を気にすることなく,ただ BFS で一番最初に $a_i $ にたどり着いたときの距離を求める答えとすることができる.偶奇が逆になっても同様である.

ただし,上記の BFS を愚直に採用すると,1つの要素の計算の時間計算量が $O(n) $ なので,全体での時間計算量が $O(n^2) $ となってしまう.

しかし,複数のノードに対して一度に計算を行うことで全体での時間計算量を $O(n) $ に抑えることが可能である.

例えば,$a_i \% 2 = 1 $ となるような要素であれば,ゴールになりうるのは $a_j \% 2 = 0 $ となるような要素であることは共通している.そのため,多始点 BFS の始点として $a_j \% 2 = 0 $ となるような要素をすべて突っ込み,$a_i \% 2 = 1 $ となる要素にたどり着いてもすべてのノードを通るまで探索を止めずに続ければ,$a_i \% 2 = 1 $ となるようなすべての要素に対してまとめて時間計算量 $O(n) $ で計算することができる.偶奇が逆になっても同様である.

以下のソースコードでは,ここからさらに (ノード番号, スタートの偶奇) で状態を区別して,すべてのノードに対して1度の BFS でまとめて計算できるようにしている.

ソースコード

 #include <iostream>
#include <vector>
#include <queue>

using namespace std;

constexpr int INF = 1 << 30;

struct State {
    int i, p, d;
    State(int i, int p, int d): i(i), p(p), d(d) {}
};

int main() {
    int n;
    cin >> n;
    vector<int> a(n), ans(n, INF);
    vector<int> rev_graph[n];
    for(int i=0;i<n;++i){
        cin >> a[i];
        if(i-a[i] >= 0) {
            rev_graph[i-a[i]].push_back(i);
        }
        if(i+a[i] < n) {
            rev_graph[i+a[i]].push_back(i);
        }
    }

    queue<State> que;
    vector<vector<bool>> used(n, vector<bool>(2, false));
    for(int i=0;i<n;++i){
        que.push(State(i, a[i] % 2, 0));   // (今いる位置,パリティ, 距離)
        used[i][a[i]%2] = true;
    }

    while(!que.empty()) {
        State s = que.front(); que.pop();
        if(a[s.i] %  2 != s.p) {
            ans[s.i] = s.d;
        }
        for(int nxt: rev_graph[s.i]) {
            if(used[nxt][s.p]) continue;
            que.push(State(nxt, s.p, s.d+1));
            used[nxt][s.p] = true;
        }
    }

    cout << ((ans[0] == INF)?-1:ans[0]);
    for(int i=1;i<n;++i) {
        cout << " " << ((ans[i] == INF)?-1:ans[i]);
    }
    cout << endl;
    
}

所感

こんなツイートしてたら E が解けなくて爆死です.謙虚にいきます…….