けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder ABC 182 D - Wandering (茶色, 400 点)

AtCoder AtCoder400点 ABC-D 茶色diff 累積和 累積max 点が移動していく問題

累積和の累積 max

問題概要

数列 A_{1}, \dots, A_{N} が与えられます。 この数列は負の要素を含むかもしれません。

数直線上の座標 0 に置かれているロボットが、以下の動作を順に行います。

  • 正の向きに A_{1} 進む。
  • 正の向きに A_{1} 進み、正の向きに A_{2} 進む。
  • 正の向きに A_{1} 進み、正の向きに A_{2} 進み、正の向きに A_{3} 進む。
  • ⋮
  • 正の向きに A_{1} 進み、正の向きに A_{2} 進み、...、正の向きに A_{N} 進む。

動作開始時から終了時までのロボットの座標の最大値を求めてください。

制約

  • 1 \le N \le 2 \times 10^{5}

考えたこと

加算する要素の個数は 1 + 2 + \dots + N = O(N^{2}) 通りあるので、愚直に足して行ったのでは間に合わない。

そこで、次のデータを持つことにしよう。

  • S_{i} := A_{1}, \dots, A_{i} の総和
  • M_{i} := S_{1}, \dots, S_{i} の最大値

S はいわゆる「累積和」で、 M は「累積和の累積 max」とでも呼ぶべきものとなっている。これらの計算は O(N) でできる。これを用いると、次のように求められる。

  • 最初の i 回の操作を終えたあとの位置を P とする
  • P に A_{1}, A_{2}, \dots, A_{i} を足していったときの
    • 結果は、 P + S_{i} となる
    • その過程での座標の最大値は、 P + M_{i} となる

以上の操作を各 i に対して実施していく。全部合わせて計算量は O(N) となる。

コード

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

int main() {
    int N;
    cin >> N;
    vector<long long> A(N), S(N+1, 0), M(N+1, 0);
    for (int i = 0; i < N; ++i) {
        cin >> A[i];
        S[i+1] = S[i] + A[i];
        M[i+1] = max(M[i], S[i+1]);
    }
    long long res = 0, cur = 0;
    for (int i = 0; i < N; ++i) {
        res = max(res, cur + M[i+1]);
        cur += S[i+1];
    }
    cout << res << endl;
}