累積和

前計算しておくことで、区間の和を $O(1)$ で求めます。添字はすべて 0-indexed で、区間は右端を含まない半開区間です。使う次元のクラスだけを貼れば動きます。

• PrefixSum(a): 長さ $N$ の列から前計算。 $O(N)$
• sum(l, r): $l \le i < r$ を満たす a[i] の和
• PrefixSum2D(a): $H \times W$ の 2 次元配列から前計算。 $O(HW)$
• sum(x1, y1, x2, y2): $x_1 \le i < x_2$, $y_1 \le j < y_2$ を満たす a[i][j] の和
• PrefixSum3D(a): $X \times Y \times Z$ の 3 次元配列から前計算。 $O(XYZ)$
• sum(x1, y1, z1, x2, y2, z2): $x_1 \le i < x_2$, $y_1 \le j < y_2$, $z_1 \le k < z_2$ を満たす a[i][j][k] の和
• 計算量: 各クエリ $O(1)$

1 次元だけなら、itertools.accumulate(a, initial=0) でも同じ表が作れます。

コード

class PrefixSum:
    def __init__(self, a):
        self.s = [0] * (len(a) + 1)
        for i, x in enumerate(a):
            self.s[i + 1] = self.s[i] + x

    def sum(self, l, r):
        # l <= i < r の a[i] の和
        return self.s[r] - self.s[l]


class PrefixSum2D:
    def __init__(self, a):
        h = len(a)
        w = len(a[0]) if h else 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] + a[i][j]
        self.s = s

    def sum(self, x1, y1, x2, y2):
        # x1 <= i < x2, y1 <= j < y2 の a[i][j] の和
        s = self.s
        return s[x2][y2] - s[x1][y2] - s[x2][y1] + s[x1][y1]


class PrefixSum3D:
    def __init__(self, a):
        nx = len(a)
        ny = len(a[0]) if nx else 0
        nz = len(a[0][0]) if ny else 0
        s = [[[0] * (nz + 1) for _ in range(ny + 1)] for _ in range(nx + 1)]
        for i in range(nx):
            for j in range(ny):
                for k in range(nz):
                    s[i + 1][j + 1][k + 1] = (
                        a[i][j][k]
                        + s[i][j + 1][k + 1] + s[i + 1][j][k + 1] + s[i + 1][j + 1][k]
                        - s[i][j][k + 1] - s[i][j + 1][k] - s[i + 1][j][k]
                        + s[i][j][k]
                    )
        self.s = s

    def sum(self, x1, y1, z1, x2, y2, z2):
        # x1 <= i < x2, y1 <= j < y2, z1 <= k < z2 の a[i][j][k] の和
        s = self.s
        return (
            s[x2][y2][z2]
            - s[x1][y2][z2] - s[x2][y1][z2] - s[x2][y2][z1]
            + s[x1][y1][z2] + s[x1][y2][z1] + s[x2][y1][z1]
            - s[x1][y1][z1]
        )