Library · Arrays · prefix sums
Prefix Sum
from itertools import accumulate
class PrefixSum:
# P[i] = a[0] + ... + a[i-1]
def __init__(self, a: list):
self.P = list(accumulate(a, initial=0))
# a[l] + ... + a[r], inclusive
def query(self, l: int, r: int):
return self.P[r + 1] - self.P[l]
class PrefixSum2D:
# P[r][c] = sum of g[0..r-1][0..c-1]
def __init__(self, g: list):
n, m = len(g), len(g[0]) if g else 0
self.P = [[0] * (m + 1) for _ in range(n + 1)] # never [[0] * (m+1)] * (n+1)
for r in range(n):
for c in range(m):
self.P[r + 1][c + 1] = g[r][c] + self.P[r][c + 1] + self.P[r + 1][c] - self.P[r][c]
# sum of g[r1..r2][c1..c2], inclusive
def query(self, r1: int, c1: int, r2: int, c2: int):
P = self.P
return P[r2 + 1][c2 + 1] - P[r1][c2 + 1] - P[r2 + 1][c1] + P[r1][c1]