はじめに
二次元配列に対して、左上と右下の座標を指定した矩形範囲の全マスに値を加える操作を何度も行い、最後に各マスの値を求める処理を考えます。操作のたびに矩形内を二重ループで走査すると、矩形が大きいほど時間がかかります。差分配列(いもす法)を使うと、1回の操作を4箇所への書き込みだけで済ませられます。
愚直な方法
高さH、幅Wのグリッドに、N回の加算操作を適用します。各操作は矩形の左上の座標(x1, y1)、右下の座標(x2, y2)、加える値vをこの順で受け取ります。矩形は半開区間[x1, x2)×[y1, y2)として、右下のマスは含みません。
grid := make([][]int, H)
for i := range grid {
grid[i] = make([]int, W)
}
for k := 0; k < N; k++ {
var x1, y1, x2, y2, v int
fmt.Fscan(in, &x1, &y1, &x2, &y2, &v)
for i := x1; i < x2; i++ {
for j := y1; j < y2; j++ {
grid[i][j] += v
}
}
}
この書き方では1回の操作につき矩形の面積ぶんのマスを書き換えます。矩形がグリッド全体に近い大きさだと、1回の操作でO(H×W)かかり、全体でO(N×H×W)です。Nが大きいと、この二重ループが処理時間の大半を占めます。
いもす法の考え方
いもす法は範囲への加算を「始点で加え、終点で打ち消す」形に置き換えます。あとで累積和を取ることで、各マスに加わった値を復元します。
一次元の場合
長さ5の配列に対して、区間[1, 4)(添字1から3)に2を加える場合を考えます。配列dの2箇所だけを書き換えます。
d[1] += 2
d[4] -= 2
書き換え後の配列dは次のとおりです。
[0, 2, 0, 0, -2]
この配列の累積和を取ります。左から順に足し上げると、d[1]で始まった+2がd[4]で打ち消されるため、区間[1, 4)に2を加えた結果になります。
[0, 2, 2, 2, 0]
二次元の場合
二次元では加算したい矩形[x1, x2)×[y1, y2)を次の4箇所への加減算で表します。
diff[x1][y1] += v
diff[x1][y2] -= v
diff[x2][y1] -= v
diff[x2][y2] += v
これは、一次元の「始点で加え、終点で打ち消す」を横方向と縦方向の両方に適用した形です。x1列から右、y1行から下のすべてのマスにvを加える操作から始め、x2列とy2行で打ち消します。二重に引いたぶんは(x2, y2)で加え戻します。
3×3のグリッドの右下2×2、つまり[1, 3)×[1, 3)に1を加える場合で確かめます。diffの初期値はH+1行、W+1列のゼロで、4箇所を書き換えます。
diff[1][1] += 1
diff[1][3] -= 1
diff[3][1] -= 1
diff[3][3] += 1
書き換え後のdiffは次のとおりです。
0 0 0 0
0 1 0 -1
0 0 0 0
0 -1 0 1
横方向、縦方向の順に累積和を取ると、3×3の範囲は次のようになります。
0 0 0
0 1 1
0 1 1
右下2×2にだけ1が加わっています。
復元は累積和
操作をすべてdiffに記録したら、累積和で各マスの最終的な値を復元します。横方向に累積和を取り、続いて縦方向に累積和を取ります。
for i := 0; i < H; i++ {
for j := 1; j < W; j++ {
diff[i][j] += diff[i][j-1]
}
}
for i := 1; i < H; i++ {
for j := 0; j < W; j++ {
diff[i][j] += diff[i-1][j]
}
}
復元後のdiff[i][j]がマス(i, j)に加えられた値の合計です。
実装
入力の読み取り、diffへの記録、復元までを通したコードは次のとおりです。diffは終点の打ち消しを記録するため、1マス余分にH+1行、W+1列で確保します。
package main
import (
"bufio"
"fmt"
"os"
)
func main() {
in := bufio.NewReader(os.Stdin)
var H, W, N int
fmt.Fscan(in, &H, &W, &N)
diff := make([][]int, H+1)
for i := range diff {
diff[i] = make([]int, W+1)
}
for k := 0; k < N; k++ {
var x1, y1, x2, y2, v int
fmt.Fscan(in, &x1, &y1, &x2, &y2, &v)
diff[x1][y1] += v
diff[x1][y2] -= v
diff[x2][y1] -= v
diff[x2][y2] += v
}
for i := 0; i < H; i++ {
for j := 1; j < W; j++ {
diff[i][j] += diff[i][j-1]
}
}
for i := 1; i < H; i++ {
for j := 0; j < W; j++ {
diff[i][j] += diff[i-1][j]
}
}
for i := 0; i < H; i++ {
for j := 0; j < W; j++ {
if j > 0 {
fmt.Print(" ")
}
fmt.Print(diff[i][j])
}
fmt.Println()
}
}
操作の記録は4箇所への加減算だけで済むため、1回の操作はO(1)です。復元の累積和はO(H×W)です。全体の計算量はO(N+H×W)となり、愚直な方法のO(N×H×W)と比べて、Nが大きいほど差が開きます。
補足
いもす法が対象にするのは加算です。加算は順番を問わず、後で累積和を取るだけで各マスの値が求まります。一方、マスへ値を上書きする操作は最後に書いた値が残るため、この方法をそのまま使えません。