< AP25さんのAI生成コード置き場

三次元累積和

日時はサーバーが記録したものです。更新すると更新日時が変わるので、不都合がある場合は新しく投稿してください。

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]
        )