Source code for titan_pylib.data_structures.cumulative_sum.cumulative_sum2D


[docs]
 1class CumulativeSum2D:
 2    """2次元累積和です。"""
 3
 4    def __init__(self, h: int, w: int, a: list[list[int]]):
 5        """``h x w`` の配列 ``a`` から2次元累積和の前計算をします。
 6        :math:`O(hw)` です。
 7
 8        Args:
 9          h (int): 行数です。
10          w (int): 列数です。
11          a (list[list[int]]):  ``CumulativeSum2D`` を構築する配列です。
12        """
13        acc = [0] * ((h + 1) * (w + 1))
14        for ij in range(h * w):
15            i, j = divmod(ij, w)
16            acc[(i + 1) * (w + 1) + j + 1] = (
17                acc[i * (w + 1) + j + 1]
18                + acc[(i + 1) * (w + 1) + j]
19                - acc[i * (w + 1) + j]
20                + a[i][j]
21            )
22        self.h = h
23        self.w = w
24        self.acc = acc
25

[docs]
26    def sum(self, h1: int, w1: int, h2: int, w2: int) -> int:
27        """長方形領域 ``[h1, h2) x [w1, w2)`` の総和を返します。
28        :math:`O(1)` です。
29        """
30        assert (
31            h1 <= h2
32        ), f"IndexError: {self.__class__.__name__}.sum({h1}, {w1}, {h2}, {w2}), h={self.h}"
33        assert (
34            w1 <= w2
35        ), f"IndexError: {self.__class__.__name__}.sum({h1}, {w1}, {h2}, {w2}), w={self.w}"
36        return (
37            self.acc[h2 * (self.w + 1) + w2]
38            - self.acc[h2 * (self.w + 1) + w1]
39            - self.acc[h1 * (self.w + 1) + w2]
40            + self.acc[h1 * (self.w + 1) + w1]
41        )