class PrefixSum3D: """ 3次元累積和 - 入力 grid は 0-index: grid[x][y][z] - sum(x1, x2, y1, y2, z1, z2) は [x1, x2) x [y1, y2) x [z1, z2) の総和を返す - 構築 O(XYZ), クエリ O(1) """ def __init__(self, grid): X = len(grid) Y = len(grid[0]) Z = len(grid[0][0]) self.X = X self.Y = Y self.Z = Z # ps[x][y][z]: # [0, x) x [0, y) x [0, z) の和 ps = [[[0] * (Z + 1) for _ in range(Y + 1)] for _ in range(X + 1)] for x in range(X): for y in range(Y): row = grid[x][y] for z in range(Z): ps[x + 1][y + 1][z + 1] = ( ps[x][y + 1][z + 1] + ps[x + 1][y][z + 1] + ps[x + 1][y + 1][z] - ps[x][y][z + 1] - ps[x][y + 1][z] - ps[x + 1][y][z] + ps[x][y][z] + row[z] ) self.ps = ps def sum(self, x1, x2, y1, y2, z1, z2): """ [x1, x2) x [y1, y2) x [z1, z2) の総和 """ ps = self.ps return ( ps[x2][y2][z2] - ps[x1][y2][z2] - ps[x2][y1][z2] - ps[x2][y2][z1] + ps[x1][y1][z2] + ps[x1][y2][z1] + ps[x2][y1][z1] - ps[x1][y1][z1] )