Overview
累積和（Prefix Sum） は、配列の部分和を高速に求めるためのテクニック。
1.累積和の基本
2.部分和を求める
3.累積和の応用
4.典型問題を解いてみる
1. 累積和の基本
整数$A_1, A_2, ..., A_N$に対し、累積和$B_i = A_1 + A_2 + ... + A_i$を考える、というときに使えるテクニック。
例 : $[3,1,4,1,5,9]$の累積和は$[3,4,8,9,14,23]$
直接計算した場合の計算量は$O(N^2)$になるが、以下のようなロジックにすると$O(N)$になる。
累積和の構築
def prefix_sum(arr):
N = len(arr)
S = [0] * (N + 1) # 0-indexed のため 1つ多く確保
for i in range(N):
S[i + 1] = S[i] + arr[i]
return S
実行イメージ
arr = [3, 1, 4, 1, 5, 9, 2]
print(prefix_sum(arr)) #出力 [0, 3, 4, 8, 9, 14, 23, 25]
実行イメージの累積和 S の計算:
| i | arr[i] | S[i+1] = S[i] + arr[i] | S |
|---|---|---|---|
| - | - | 初期値 | [0] |
| 0 | 3 | S[1] = 0 + 3 | [0, 3] |
| 1 | 1 | S[2] = 3 + 1 | [0, 3, 4] |
| 2 | 4 | S[3] = 4 + 4 | [0, 3, 4, 8] |
| 3 | 1 | S[4] = 8 + 1 | [0, 3, 4, 8, 9] |
| 4 | 5 | S[5] = 9 + 5 | [0, 3, 4, 8, 9, 14] |
| 5 | 9 | S[6] = 14 + 9 | [0, 3, 4, 8, 9, 14, 23] |
| 6 | 2 | S[7] = 23 + 2 | [0, 3, 4, 8, 9, 14, 23, 25] |
2. 部分和を求める
部分和、つまり区間[l,r]の合計を求めるのは、$S[r+1]−S[l]$。
A = [3, 1, 4, 1, 5, 9, 2]
S = prefix_sum(A) #1.累積和の構築の関数を呼び出す
l, r = 2, 5 # 0-indexed
print(S[r+1] - S[l]) # 出力: 4 + 1 + 5 + 9 = 19
・従来の計算だと $O(N)$ かかるが、累積和なら$ O(1)$
・累積和を事前計算することで高速に処理できる
3. 累積和の応用
(1) 区間の平均を求める
部分和がすぐに求まるので、平均も O(1) で求められる。
def range_average(S, l, r):
return (S[r+1] - S[l]) / (r - l + 1)
A = [3, 1, 4, 1, 5, 9, 2]
S = prefix_sum(A) #1.累積和の構築の関数を呼び出す
l, r = 2, 5
print(range_average(S, l, r)) # 平均値 (4+1+5+9)/4 = 4.75
(2) 2次元累積和
問題: グリッド N×M の 長方形領域の合計を O(1) で求める。
def build_2d_prefix_sum(grid):
H, W = len(grid), len(grid[0])
S = [[0] * (W + 1) for _ in range(H + 1)]
for i in range(H):
for j in range(W):
S[i+1][j+1] = S[i][j+1] + S[i+1][j] - S[i][j] + grid[i][j]
return S
def query_2d(S, x1, y1, x2, y2):
return S[x2+1][y2+1] - S[x1][y2+1] - S[x2+1][y1] + S[x1][y1]
grid = [
[3, 1, 4],
[1, 5, 9],
[2, 6, 5]
]
S = build_2d_prefix_sum(grid)
print(query_2d(S, 1, 1, 2, 2)) # 出力: 5 + 9 + 6 + 5 = 25
・2D累積和を使うと、部分領域の合計が$O(1)$
・通常の計算では $O(NM)$ かかるが、高速化可能
4. 典型問題を解いてみる
例題1. 1次元累積和（区間和）
問題: N 個の整数 A と Q 個のクエリ l, r が与えられる。各クエリの区間 [l, r] の和を求めよ。
回答
N, Q = map(int, input().split())
A = list(map(int, input().split()))
S = prefix_sum(A)
for _ in range(Q):
l, r = map(int, input().split())
print(S[r+1] - S[l])
例題2. 累積和を使った平均計算
問題: N 個の整数 A が与えられる。Q 個のクエリに対して区間 [l, r] の平均を求めよ。
回答
N, Q = map(int, input().split())
A = list(map(int, input().split()))
S = prefix_sum(A)
for _ in range(Q):
l, r = map(int, input().split())
print((S[r+1] - S[l]) / (r - l + 1))
例題3. 2次元累積和（長方形領域の合計）
問題: H×W のグリッドが与えられる。Q 個のクエリ (x1, y1, x2, y2) に対し、指定領域の合計を求めよ。
回答
H, W, Q = map(int, input().split())
grid = [list(map(int, input().split())) for _ in range(H)]
S = build_2d_prefix_sum(grid)
for _ in range(Q):
x1, y1, x2, y2 = map(int, input().split())
print(query_2d(S, x1, y1, x2, y2))
例題4. 条件付き累積和
問題: N 個の整数 A から、偶数のみの累積和を求め、クエリに答えよ。
回答
N, Q = map(int, input().split())
A = list(map(int, input().split()))
S = [0] * (N+1)
for i in range(N):
S[i+1] = S[i] + (A[i] if A[i] % 2 == 0 else 0)
for _ in range(Q):
l, r = map(int, input().split())
print(S[r+1] - S[l])