競プロで度々使う累積和の使い方
1次元の累積和
以下の問題を考える。
連続するk個の総和の最大値を求める問題。
シンプルに考えると、連続するk個の値はn-k個のパターンがあり、それぞれを求めて最大値を求めていく。
これだと、連続するk個の総和を求める処理が計算量にするとO(N2)かかる。
コードで表現すると以下のような感じ。
package main
import (
"bufio"
"fmt"
"os"
"strconv"
)
var scanner = bufio.NewScanner(os.Stdin)
func nextInt() int {
scanner.Scan() // トークンの読み込み
i, e := strconv.Atoi(scanner.Text()) // 読み込んだトークンをTextメソッドで取り出し
if e != nil {
panic(e)
}
return i
}
func main() {
scanner.Split(bufio.ScanWords) // 空白区切りで受け取るので引数にbufio.ScanwWordsを渡す
N := nextInt()
A := make([]int, N)
for i := 0; i < N; i++ {
A[i] = nextInt()
}
for k := 0; k < N; k++ {
var max int
for i := 0; i < N-k; i++ {
a := A[i : i+k+1]
var sum int
for j := range a {
sum += a[j]
}
if max < sum {
max = sum
}
}
fmt.Println(max)
}
}
これだと全体の計算量はO(N3)になってしまい、TLEする。
こういう時は累積和を考える。
累積和Sは、S[n+1] = S[n] + a[n+1] と定義される。
つまり、累積和S[n+1] は 配列a[0] から a[n+1] までの累積の和である。（数IIBでやったよね懐かしい）
この累積和をあらかじめ求めておくと、連続するk個の和は S[i+k] - S[i] で求められるので、連続するk個の総和を求める処理がO(N)、全体でO(N2)となりACできる。
コードで表現すると以下になる。
package main
import (
"fmt"
)
func main() {
var N int
fmt.Scan(&N)
A := make([]int, N)
for i := 0; i < N; i++ {
fmt.Scan(&A[i])
}
s := make([]int, N+1) // 累積和
for i := 0; i < N; i++ {
s[i+1] = s[i] + A[i]
}
for k := 1; k < N+1; k++ { // 連続するk個の区画について考える
var max int
for i := 0; i < N+1; i++ {
if i+k > N {
continue
}
if max < s[i+k]-s[i] { // 連続するk個の和
max = s[i+k] - s[i]
}
}
fmt.Println(max)
}
}
このように特定の範囲での総和を求める問題などでは、あらかじめ累積和を求めておくことで計算量を下げられることがよくある。
2次元の累積和
次に以下の問題を考える。
この問題では長方形になるようなブロックを列挙して面積の最大値を求める。
普通に長方形になるようなブロックを選ぼうとすると、長方形かどうかの判定も書かねばならないので煩雑になりそう。
長方形になるようなブロックの選び方を列挙しなければいけないときは2次元の累積和を使うと比較的楽に実装できる。
2次元累積和の考え方
2次元の累積和では、(0, 0) から配列の座標までの累積和を記録する。
つまりこの問題で言うと、2次元累積和S[i][j] では (0, 0) から (i, j) までのブロックの価値の累積和を記録する。
この際、2次元累積和は S[i+1][j+1] = S[i+1][j] + S[i][j+1] - S[i][j] + A[i+1][j+1] と定義される（ブロックの足し引き考えればイメージできる）
このように、あらかじめ2次元累積和で長方形の形で価値の累積和を記録しておくと、任意の長方形の価値を抽出する際も長方形から任意の長方形を足し引きすることで求められる。
(k, l) から (i, j) からなる長方形は S[i][j] - S[k][j] - S[i][l] + S[k][l] と表せる。
コードで表現すると、以下のようになる。
package main
import (
"fmt"
)
func main() {
var H, W, K, V int
fmt.Scan(&H, &W, &K, &V)
A := make([][]int, H)
sum := make([][]int, H+1) // 累積和: (0,0)と(i,j)を頂点とする長方形の地価の合計値
for i := range sum {
sum[i] = make([]int, W+1)
}
for i := 0; i < H; i++ {
A[i] = make([]int, W)
for j := 0; j < W; j++ {
fmt.Scan(&A[i][j])
sum[i+1][j+1] = sum[i][j+1] + sum[i+1][j] - sum[i][j] + A[i][j] // sum[i][j]はsum[i][j+1]とsum[i+1][j]の重複部分
}
}
var max int
for i := 0; i < W+1; i++ { // i, jで大元の長方形を決定
for j := 0; j < H+1; j++ {
for k := 0; k < i; k++ { // k, lで切り出す長方形を決定
for l := 0; l < j; l++ {
area := (i - k) * (j - l) // 面積
areaCost := sum[j][i] - sum[j][k] - sum[l][i] + sum[l][k]
buildingCost := area * K
cost := areaCost + buildingCost
if V < cost {
continue
}
if max < area {
max = area
}
}
}
}
}
fmt.Println(max)
}
このように2次元累積和は長方形で物事を考えなくてはいけない場合に、計算量を下げたり計算が容易になったりする。
