ABC 253

Table of Contents

E - Distance Sequence

https://atcoder.jp/contests/abc253/tasks/abc253_e

$f(i,j)$ を $A_i = j$ となる場合の数とすると

\begin{align} f(1,j) &= 1 ~~~(\forall j \in [1,M]) \nonumber \\ f(i+1,j) &= \sum_{k=1}^{j-K} f(i,k) + \sum_{k=j+K}^{M} f(i,k) \nonumber \end{align}

となる。naive に実装すると以下のようになり、計算量が $O(NM^2)$ となるため問題の制約では TLE になる。

vector<vector<mint>> dp(N + 1, vector<mint>(M + 1, 0));
rep(i, M) dp[1][i + 1] = 1;
rep2(i, 2, N + 1) {
    rep2(k, 1, M + 1) {
        rep2(j, 1, M + 1) {
            if (abs(k - j) >= K)
                dp[i][k] += dp[i - 1][j];
        }
    }
}

mint ans = 0;
rep2(i, 1, M + 1) ans += dp[N][i];
cout << ans.val() << endl;

ここで dp[i-1] の累積和計算しておくことで一番内側の loop 内の処理を $O(1)$ で行うことができる。 $K = 0$ のときに dp[i-1][j] からの寄与をダブルカウントしないように注意する

void solve() {
    ll N, M, K;
    cin >> N >> M >> K;

    vector<vector<mint>> dp(N + 1, vector<mint>(M + 1, 0));
    rep(i, M) dp[1][i + 1] = 1;
    rep2(i, 2, N + 1) {
        vector<mint> cumsum(M + 2, 0);

        rep2(j, 1, M + 1) {
            cumsum[j + 1] += cumsum[j] + dp[i - 1][j];
        }

        ll offset = 1;
        if (K == 0)
            offset = 0;

        rep2(j, 1, M + 1) {
            dp[i][j] += cumsum[M + 1] - cumsum[min(M + 1, j + K)];
            dp[i][j] += cumsum[max(1ll, j - K + offset)] - cumsum[1];
        }
    }

    mint ans = 0;
    rep2(i, 1, M + 1) ans += dp[N][i];
    cout << ans.val() << endl;
}