K Integers Summations [第六回 アルゴリズム実技検定 D]

https://atcoder.jp/contests/past202104-open/tasks/past202104_d

前提知識

• 累積和

解説

https://atcoder.jp/contests/past202104-open/submissions/22645520

ここから点数があがり、単なる実装問題ではなくなってくる。
計算量を意識していこう。

愚直解法

全てのxについて、A[x]からA[x+K-1]までの和を取っていくような実装をする。
xはN-K+1通りあり、そこからK個分の総和を取るループを作るので、(N-K+1)×K回和を取る処理が実行されるので、O(NK)となる。
これは今回の制約では10¹⁰を超えてしまうので間に合わない。
一般に計算量は、色々な要因があって揺らぎがあるが、10⁷くらいが限界である。

間に合わないが、実装しておいた。
まずは愚直解法が分からないと、それを最適化した解法も分からないだろうと思う。
上の説明がピンとこない場合は、↓の実装を参考にしてみてほしい。
https://atcoder.jp/contests/past202104-open/submissions/22645442

間に合わないのだが、愚直解法が適切な解法の手助けになることもある。
xの全探索は行うのだが、A[x]からA[x+K-1]までの和を高速に取る方法はないだろうか。
実は累積和という手法を使うと、O(1)で区間の総和を取ることができ、O(N)に計算量が改善される。

累積和

累積和の詳しい説明はググって詳解しているページで見てみてほしいが、軽く以下に説明しておく。
事前に以下のような配列を用意しておく。

tot[x] = A[0] + A[1] + ... + A[x]

これは以下のような漸化式を使えば高速に構築できる。

tot[0] = A[0]
tot[i] = tot[i - 1] + A[i]

こうすると、今回求めたい区間和は以下のように書ける。

A[x]からA[x+K-1]までの和 = tot[x+K-1] - tot[x-1]

式変形をすればそうなることが分かる。これで計算量がO(K)からO(1)に改善されるので全体がO(N)となり間に合う。

実装上の注意点

以下2つある。

• totの中身はintで受け取れる範囲を超えてしまうので、long longで取ること
• 問題に「出力する行数が大きくなる場合があるので、注意せよ。」とあるが、C++ならprintfで出力しないと間に合わないかもしれない
  • coutでも何かすれば早くなるんだったかな？忘れた

int N, K, A[501010];
ll tot[501010];
//---------------------------------------------------------------------------------------------------
void _main() {
    cin >> N >> K;
    rep(i, 0, N) cin >> A[i];

    tot[0] = A[0];
    rep(i, 1, N) tot[i] = tot[i - 1] + A[i];

    rep(x, 0, N - K + 1) {
        ll ans = tot[x + K - 1];
        if (0 < x) ans -= tot[x - 1];
        printf("%lld\n", ans);
    }
}