累積和

読了目安 約3分

先頭からの合計を前計算し、区間和クエリに O(1) で答える。定石は prefix[i + 1] = prefix[i] + a[i]。

「区間の合計」を何度も聞かれる問題では、毎回ループで足すと間に合いません。 先頭からの合計を前もって作っておけば、どの区間の合計も引き算 1 回で出せます。 この前計算した配列を累積和と呼びます。

前計算の定石

長さ N の配列 a に対し、長さ N + 1 の配列 prefix を作ります。 prefix[i] は「先頭から i 個の合計」です。 埋め方の定石は prefix[i + 1] = prefix[i] + a[i] です。

let a = [3, 1, 4, 1, 5, 9]

var prefix = Array(repeating: 0, count: a.count + 1)
for i in 0..<a.count {
    prefix[i + 1] = prefix[i] + a[i]
}
print(prefix)

prefix[0] は「0 個の合計」なので 0 です。 長さを 1 つ余らせるおかげで、次の引き算がきれいに書けます。

区間和は引き算 1 回

a[l] から a[r - 1] までの合計は prefix[r] - prefix[l] です。 「先頭から r 個の合計」から「先頭から l 個の合計」を取り除くと、間の部分だけが残るからです。

let a = [3, 1, 4, 1, 5, 9]
var prefix = Array(repeating: 0, count: a.count + 1)
for i in 0..<a.count {
    prefix[i + 1] = prefix[i] + a[i]
}

// a[1] + a[2] + a[3] = 1 + 4 + 1
print(prefix[4] - prefix[1])
// 全体の合計
print(prefix[6] - prefix[0])

試してみよう: a[2] から a[5] までの合計 19 を、引き算 1 回で出力してみてください。

計算量

前計算が O(N)、クエリ 1 個への回答が O(1) です。 クエリが Q 個あっても、全体で O(N + Q) です。

毎回足す愚直な方法は O(NQ) です。 N と Q が 10 万なら 100 億回になり、間に合いません。

いもす法

逆に「区間への足し算」が何度も来る問題には、いもす法が使えます。 足す量を区間の両端にメモしておき、最後に累積和を 1 回取ると、全部の足し算がまとめて反映されます。

var diff = Array(repeating: 0, count: 7)
diff[1] += 10   // a[1] から a[3] に 10 を足す
diff[4] -= 10   // a[4] からは足さない
diff[2] += 1    // a[2] から a[5] に 1 を足す
diff[6] -= 1

var value = 0
for i in 0..<6 {
    value += diff[i]
    print("a[\(i)] への加算: \(value)")
}

メモは区間 1 個につき 2 か所だけなので、Q 回の区間更新も O(N + Q) で処理できます。

学びどころ

演習

1 行目に整数 N と Q が入力されます。 2 行目に N 個の整数 a[0] … a[N - 1]、続く Q 行に整数 l と r が入力されます。 各行について、a[l] から a[r - 1] までの合計を出力してください。

制約: 1 ≤ N ≤ 100000、1 ≤ Q ≤ 100000、0 ≤ a[i] ≤ 100、0 ≤ l < r ≤ N。 クエリごとにループで足すと最大 100 億回の足し算になり、間に合いません。 累積和を前計算してから答えてください。

入力例:
5 3
1 2 3 4 5
0 5
1 4
2 3

出力例:
15
9
3

let firstLine = readLine()!.split(separator: " ").map { Int($0)! }
let n = firstLine[0]
let q = firstLine[1]
let a = readLine()!.split(separator: " ").map { Int($0)! }

var prefix = Array(repeating: 0, count: n + 1)
for i in 0..<n {
    prefix[i + 1] = prefix[i] + a[i]
}

for _ in 0..<q {
    let lr = readLine()!.split(separator: " ").map { Int($0)! }
    print(prefix[lr[1]] - prefix[lr[0]])
}