累積和を使った計算量の削減について【競技プログラミング】

# はじめに
- 先日、[dp(動的計画法)の投稿](https://qiita.com/kei_tnk/items/9a7380b3fa954688f254)をしたけど、コッテリしていて若干消化不良気味で気持ち悪いので、あっさりした「累積和」について投稿して、体調を整えたい。
- ちなみに、swiftで書きます。

# 累積和について
- 計算量削減テクニックの一つで、AtCoderのABCコンテストなら、c問題当たりで出そうな感じかな。
- 累積和の活躍の最も分かり易い例は、配列の連続区間の合計値を出す作業（クエリ）を繰り返し行う場面かな。
- 例えば、下記のような入力があるとします。
```
5 2 // 配列サイズ5,クエリ数2
2 4 7 1 2 // 配列（サイズ5）
1 3 // クエリ条件① 配列の1つ目から3つ目までの値を合計せよ ⇒出力：13
2 5 // クエリ条件② 配列の2つ目から5つ目までの値を合計せよ ⇒出力：14
```
- 勿論、上記を求めるコードは、深く考えずに書くことが出来ます。
```swift
let (n,q) = [readLine()!.split(separator:" ").map{Int($0)!}].map{($0[0],$0[1])}[0]
let vs = readLine()!.split(separator:" ").map{Int($0)!}

for _ in 0..<q {
    let (l,r) = [readLine()!.split(separator:" ").map{Int($0)!}].map{($0[0]-1,$0[1]-1)}[0] // 1ずらす

    /// クエリ計算：計算量O(n) -- (※1)    
    var sum = 0
    for i in l...r {
        sum += vs[i]
    }
    
    print(sum)
}
```
- ちゃんと、答えは①13，②14となったよ。
- でも、このアルゴリズムの計算量はO(q･n)だから、$n=10^6,q=10^4$のような規模だと、計算量O($10^{10}$)になって、TLE(実行時間制限超過)になっちゃう。
- これを防ぐのが、累積和を使ったテクニックです。累積和を使うと、上記(※1)のクエリ計算の計算量O(n)を大幅に削減できます。
- 元の配列vsに対し、累積和の配列ssの各要素は以下のように作成します。
    - ss[i] = Σvs(j) {j in 0..<i}
    - つまり、累積和ssのインデックスiの要素の値は、元の配列のインデックスiより前の値の合計です。
        - 元の配列のインデックスiの値は「含まない」ことに注意！
        - 累積和ssのサイズは「n+1」であり、元の配列vsのサイズ「n」より1つ大きいことに注意！
- このssを使うと、クエリ計算は下記のように書き換えられます。
```swift
    /// クエリ計算：計算量O(n) -- (※2)
    let sum = ss[r+1] - ss[l] // ss[r+1]はvs[0]〜vs[r]の合計、ss[l]はvs[0]〜vs[l-1]の合計
```
- 上記(*2)の計算量はO(1)となります。ただし、ssを準備するための計算量O(n)が発生するので、全体の計算量は、O(n+q)となりますが、元の計算量O(q･n)より大幅に削減されます。$n=10^6,q=10^4$の場合、累積和を使うと、計算量O($10^6$)になって、余裕ですね。
- 累積和を利用したコードを書くと、こんな感じ。
```swift
let (n,q) = [readLine()!.split(separator:" ").map{Int($0)!}].map{($0[0],$0[1])}[0]
let vs = readLine()!.split(separator:" ").map{Int($0)!}

var ss:[Int] = [0] // これが重要！！！
var s = 0
for v in vs {
    s += v
    ss.append(s)
}

for _ in 0..<q {
    let (l,r) = [readLine()!.split(separator:" ").map{Int($0)!}].map{($0[0]-1,$0[1]-1)}[0] // 1ずらす
    
    let sum = ss[r+1] - ss[l]
    
    print(sum)
}
```
- 上記コードにも書いているけど、ssの最初の要素は必ず[0]が入る。また、この[0]により、ssのサイズはvsのサイズより1つ大きくなる。

# 実際のコンテスト問題①　（引っかけ問題）
- [ABC122のc問題](https://atcoder.jp/contests/abc122/tasks/abc122_c)が例として良いと思う。
```
8 3 // 文字数8、クエリ数3
ACACTACG // "AC"を含む文字列
3 7 // クエリ① 3〜7文字目の間に"AC"を何文字含むか？
2 3 // クエリ② 2〜3文字目の間に"AC"を何文字含むか？
1 8
```
- 上記を解くコードはこんな感じ
```swift
let (n,q) = [readLine()!.split(separator:" ").map{Int($0)!}].map{($0[0],$0[1])}[0]

//文字列から各文字配列化 
var cs = Array(readLine()!).map{String($0)} // 要素はString型(重い)

var ss:[Int] = [0]
var s = 0
var old = ""
for c in cs {
    if old + c == "AC" { s += 1 }
    ss.append(s)
    old = c
}

for _ in 0..<q {
    let (l,r) = [readLine()!.split(separator:" ").map{Int($0)!}].map{($0[0]-1,$0[1]-1)}[0] // 1ずらす
    
    var sum = ss[r+1] - ss[l]
    if l>0 && cs[l-1] + cs[l] == "AC" {sum -= 1} // これが重要！！！
    
    print(sum)
}
```
- 上記コードで、累積和の配列ssを作成する部分は、特に問題なく出来ると思う。
- この問題のキモは「これが重要！！！」と書いている部分。これがないと、入力例のクエリ②で結果が1となってしまう。
- クエリ②を見ると、2〜3文字目("CA")の中に"AC"はゼロだから、結果は0とすべきである。しかし、累積和を使った解法では、$sum = ss[r+1]-ss[l]$としているから、ss=[0,<font color=red>0</font>,<font color=green>1</font>,1,2,2,2,3,3]における、ss[1]=0とss[2]=1の差から1となってしまった。
- つまり、先頭部分の"AC"のお尻の"C"のみで1カウントされてしまうのを戻すことをしている。
- これに気付かないと解けないけど、この問題は親切なことに、入力例1で引っかけ部分に気付かせてくれている。

# 実際のコンテスト問題②　（３次元配列）
- 次は、3次元配列を使った問題である、[ABC366のd問題](https://atcoder.jp/contests/abc366/tasks/abc366_d)を例にする
```
3 // 1辺のサイズが3の三次元配列
733 857 714 // x=1 y=1 z=1,2,3
956 208 257 //     y=2 z=1,2,3
123 719 648 //     y=3 z=1,2,3
840 881 245 // x=2
245 112 746
306 942 694
58 870 849 // x=3
13 208 789
687 906 783
8 // クエリ8
3 3 3 3 1 1 //クエリ① x:3~3、y:3~3、z:1~1 の区間の三次元配列の要素合計はいくらか？
1 3 2 3 3 3 //クエリ② x:1~3、y:2~3、z:3~3　の区間の三次元配列の要素合計はいくらか？
2 2 2 3 1 1
1 3 1 1 1 1
2 3 2 3 2 3
1 2 1 1 1 2
3 3 2 2 1 3
1 2 2 3 2 3
```
- とりあえず、三次元配列vsssを取得するコードは次のとおり
```swift
//整数入力
let N = Int(readLine()!)!
//3次元配列
var A:[[[Int]]] = []
for _ in 0..<N {
    var vss:[[Int]] = []
    for _ in 0..<N {
        let vs = readLine()!.split(separator:" ").map{Int($0)!}
        vss.append(vs)
    }
    A.append(vss)
}
```
- 上記3次元配列の各要素へは、A[x][y][z] {x:0...(N-1),y:0...(N-1),z:0...(N-1)}となる。
- さて、ここで、累積和を格納する配列をSとした場合、どのように作れば良いだろうか？
- 1次元の配列の時を踏襲して、
    - S[i][j][k] = ΣΣΣ A[x][y][z] {x:0..<i , y:0..<j , z:0..<k} -- (※)
- しかし、このSをAの要素毎に加算していてはTLEになってしまうので、下記の特性を用いて、計算量を削減する。
    - $S[i][j][k]$は、下記の合計となる。
        ① $S[i-1][j-1][k-1]$
        ② $A[i-1][j-1][k-1]$
        ③ $S[i-1][j][k] + S[i][j-1][k] + S[i][j][k-1]$
        ④ $-S[i][j-1][k-1] - S[i-1][j][k-1] - S[i-1][j-1][k]$
    - ①と②が含まれるのは感覚的に分かるだろうね。
    - ③と④については、ちょっと分かりづらいけど、そういうモノだと記憶して欲しい。
    - また、i,j,kのどれか一つでも0だと、上記の定義上(※)による計算ができないが、1次元の配列の時と同様、$S[i][j][k]はゼロとなるものとする。
- よって、累積和の配列は下記コードで算出できる。
```swift
var S = [[[Int]]](repeating:[[Int]](repeating:[Int](repeating:0,count:N+1),count:N+1),count:N+1)
for i in 0...N {
    for j in 0...N {
        for k in 0...N {
            if i * j * k == 0 { S[i][j][k] = 0 } else {
            S[i][j][k] = S[i-1][j-1][k-1]
                        + A[i-1][j-1][k-1]
                        + S[i-1][j][k] + S[i][j-1][k] + S[i][j][k-1] 
                        - S[i][j-1][k-1] - S[i-1][j][k-1] - S[i-1][j-1][k]
            }
        }
    }
}
```
- さて、最後にクエリの結果をどのように算出すれば良いだろうか？
- 頭の中で立方体の形を考えれば、まあ、次の形になることは分かるだろうか。
```swift
let Q = Int(readLine()!)!

for _ in 0..<Q {
    let vs = readLine()!.split(separator:" ").map{Int($0)! - 1} // -1ずらす
    let (lx,rx,ly,ry,lz,rz) = (vs[0],vs[1],vs[2],vs[3],vs[4],vs[5])
    
    let ans = S[rx+1][ry+1][rz+1] - S[lx][ly][lz]
            + S[rx+1][ly][lz] + S[lx][ry+1][lz] + S[lx][ly][rz+1]
            - S[lx][ry+1][rz+1] - S[rx+1][ly][rz+1] - S[rx+1][ry+1][lz]
    
    print(ans)
    
}
```
- 提出したら600ms程度でACできた。

# おわりに
- 累積和はあまり複雑じゃ無いから口直しに良いかな〜と思ったけど、三次元配列に手を出したら、結局頭が混乱しちゃったよ。
- でも、これでもd問題なのよね。どうやら、夏の間に4完するのは無理ぽ...

【追加】累積和の考えを利用して、二重forループや三重forループの内部ループを分離して計算量を削減する「いもす法」を投稿したから、そちらも参考にして欲しい。

https://qiita.com/kei_tnk/items/79ca53a0a84054ba09a0