[自分用] 競プロのための累積和入門

5/11に開催されたABCコンテストにおいて、C問題が何度もTLEとなり解くことが出来なかった。
その対策として、「累積和」を取り入れる必要があると思ったので、学習しながらここに記していく。

参考にした記事

累積和とは

・・配列内の総和を求める処理を爆速にする手法

累積和の考え方

例えば、以下のような要件があったとする。

要件：以下の10個の自然数のうち、連続する3つの和の最大値を求めよ

1,8,4,2,5,6,3,7,1,4

手順1

配列Aに10個の自然数を格納。
配列A[1,8,4,2,5,6,3,7,1,4]

手順2

配列Bの0+N番目に、配列Aの0からN-1番目の和を格納。
配列B[0,1,9,13,15,20,26,29,36,37,41]

上記からわかることは以下。

配列Aのiからx番目までの和は、配列Bのx+1とiの差に一致する。
例えば、配列Aの1番目から3番目までの和を配列Bから求めるとする。

i=1,x=3のため、
配列Aは
配列A[1] + 配列A[2] + 配列A[3] = 8+4+2 = 14
配列Bは
配列B[4] - 配列B[1] = 15-1 = 14

このように、累積和を使えば、連続するn個の和を高速で算出することが出来る。

これを利用して最初の問題を解くと、「連続する3つの和」は配列B[n+3] - 配列B[n]に一致するため、連続する3つの和の最大値のみを常に保持しておけば、配列Bの末尾まで探索が完了した時に保持している最大値が配列Aにおける連続する3つの和の最大値ということになる。
この問題の場合、配列Bの36から20を引いた16が連続する3つの和の最大値となる。

配列B[0,1,9,13,15,20,26,29,36,37,41]

累積和をJavaで実装

具体的にJavaでプログラムを書いて問題を解いてみる。

public static void main(String[] args){

  int[] A = {1,8,4,2,5,6,3,7,1,4};

  int[] B = new int[A.length+1];

  B[0] = 0;

  for(int i=1;i<A.length+1;i++){
    B[i]= B[i-1] + A[i-1]; //配列Bに累積和を格納
  }

  int maxNum = 0;
  for(int i=0;i+3<A.length;i++){
    if(B[i+3]-B[i]>maxNum){
      maxNum = B[i+3]-B[i]; 
      //現状保持している連続する3つの和の最大値より大きかった場合、最大値を変更
    }
  }

  System.out.println(maxNum); //最大値を出力

}

コンソールには16が出力される。

まとめ

累積和を使うことで、2重ループなどを行う必要がなく、処理が高速化する。競プロで勝ち抜くために必要な知識だと感じたので、記事にした。
今後も必要に応じて記事を更新していくので、Qiita、Twitterのフォロー推奨。