はじめに
- 先日、dp(動的計画法)の投稿をしたけど、コッテリしていて若干消化不良気味で気持ち悪いので、あっさりした「累積和」について投稿して、体調を整えたい。
- ちなみに、swiftで書きます。
累積和について
- 計算量削減テクニックの一つで、AtCoderのABCコンテストなら、c問題当たりで出そうな感じかな。
- 累積和の活躍の最も分かり易い例は、配列の連続区間の合計値を出す作業（クエリ）を繰り返し行う場面かな。
- 例えば、下記のような入力があるとします。
5 2 // 配列サイズ5,クエリ数2
2 4 7 1 2 // 配列（サイズ5）
1 3 // クエリ条件① 配列の1つ目から3つ目までの値を合計せよ ⇒出力：13
2 5 // クエリ条件② 配列の2つ目から5つ目までの値を合計せよ ⇒出力：14
- 勿論、上記を求めるコードは、深く考えずに書くことが出来ます。
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を使うと、クエリ計算は下記のように書き換えられます。
/// クエリ計算：計算量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$)になって、余裕ですね。
- 累積和を利用したコードを書くと、こんな感じ。
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問題が例として良いと思う。
8 3 // 文字数8、クエリ数3
ACACTACG // "AC"を含む文字列
3 7 // クエリ① 3〜7文字目の間に"AC"を何文字含むか？
2 3 // クエリ② 2〜3文字目の間に"AC"を何文字含むか？
1 8
- 上記を解くコードはこんな感じ
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,0,1,1,2,2,2,3,3]における、ss[1]=0とss[2]=1の差から1となってしまった。
- つまり、先頭部分の"AC"のお尻の"C"のみで1カウントされてしまうのを戻すことをしている。
- これに気付かないと解けないけど、この問題は親切なことに、入力例1で引っかけ部分に気付かせてくれている。
実際のコンテスト問題② （３次元配列）
- 次は、3次元配列を使った問題である、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を取得するコードは次のとおり
//整数入力
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]はゼロとなるものとする。
- $S[i][j][k]$は、下記の合計となる。
- よって、累積和の配列は下記コードで算出できる。
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]
}
}
}
}
- さて、最後にクエリの結果をどのように算出すれば良いだろうか？
- 頭の中で立方体の形を考えれば、まあ、次の形になることは分かるだろうか。
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ループの内部ループを分離して計算量を削減する「いもす法」を投稿したから、そちらも参考にして欲しい。