【AtCoder】ABC 401 C - K-bonacci | 緑コーダーが解くAtCoder

atcoder.jp

配点: 300 点 / 実行時間制限: 2 sec / メモリ制限: 1024 MB / Difficulty: ???

問題概要
制約
考察
実装例
コメント

問題概要

正整数 $N,K$ が与えられる。長さ $N+1$ の数列 $A$ の各要素の値を、以下の方法で定義する。

$0\leq i \lt K$ のとき、 $A_i=1$
$K\leq i$ のとき、 $A_i=A_{i-K}+A_{i-K+1}+\cdots+A_{i-1}$

このとき、 $A_N$ を $10^9$ で割った余りを求めよ。

制約

$1\leq N, K\leq 10^{6}$
入力は全て整数。

考察

愚直に $A_i$ を求めるやり方だと、 $K \leq i$ のときに毎回 $K$ 回の計算が発生することになり、全体の計算量は $O(NK)$ となってしまう。

そこで、この部分の処理を累積和を使って高速化しよう。

今回は、 $A_i$ の値が求まった後に累積和テーブルが順次定まっていくことになる。

また、本問では $N, K$ の大小関係には触れられていないことに注意したい。

任意の値での剰余を取るには、 ACL の modintを使うと便利である。

実装例

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

#if __has_include(<atcoder/all>)
#include <atcoder/all>
using namespace atcoder;
#endif

#define rep(i, start, end) for (auto i = (start); (i) < (end); (i)++)

// ======================================== //

using mint = static_modint<1000000000>;

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

    vector<mint> A(N + 1), S(N + 1);
    rep(i, 0, min(K, N + 1))
    {
        A[i] = 1;
        S[i] = A[i];
        if (i >= 1)
        {
            S[i] += S[i - 1];
        }
    }

    rep(i, K, N + 1)
    {
        A[i] = S[i - 1];
        if (i - K >= 1)
        {
            A[i] -= S[i - K - 1];
        }

        S[i] = S[i - 1] + A[i];
    }

    cout << A[N].val() << endl;

    return 0;
}