AtCoder ABC 322 C - Festival (灰色, 300 点)

for 文を回していくときに「前回の値を活用する」というテクが登場する問題。このテクは、後に「累積和」や「DP」を学ぶ際の大切な基礎となる。

• 問題へのリンク

問題概要

日間のうち、 日目には花火が上がる。

各日について、何日後に最初の花火が上がるかを答えよ。

制約

• 

考えたこと

二重ループを書いて TLE にならないなら簡単。実際には TLE になるので少し工夫する。

まず、 日間を マスに見立てることにする。そして、花火のない日は白色マス、花火のある日は黒色マスとする。なお、このように、 日間という設定を マスという設定にして可視化するのは、地味ながらも結構使えるテク。

そして、下図のように、後ろから見ていって

• 白色マスの場合は、その後ろのマスの値 + 1 とする
• 黒色マスの場合は、0 とする

とすれば OK。

最後に計算量を見積もる。for 文を 1 回回すだけなので、 の計算量となる。

コード

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

int main() {
    int N, M;
    cin >> N >> M;
    
    // N マスを作る
    vector<bool> is_hanabi(N, 0);
    for (int i = 0; i < M; ++i) {
        int A;
        cin >> A;
        --A;
        is_hanabi[A] = true;
    }
    
    // 答えを用意
    vector<int> res(N);
    
    // N 日目は花火と決まっている
    res[N-1] = 0;
    
    // 後ろから見ていく
    for (int i = N-2; i >= 0; --i) {
        if (!is_hanabi[i]) {
            // 花火でない日 (白色マス) は前回の値 + 1
            res[i] = res[i+1] + 1;
        } else {
            // 花火の日 (黒色マス) は 0 にリセット
            res[i] = 0;
        }
    }
    
    // 出力
    for (int i = 0; i < N; ++i) cout << res[i] << endl;
}