cumulative_sum2D

ソースコード

from titan_pylib.data_structures.cumulative_sum.cumulative_sum2D import CumulativeSum2D

view on github

展開済みコード

 1# from titan_pylib.data_structures.cumulative_sum.cumulative_sum2D import CumulativeSum2D
 2class CumulativeSum2D:
 3    """2次元累積和です。"""
 4
 5    def __init__(self, h: int, w: int, a: list[list[int]]):
 6        """``h x w`` の配列 ``a`` から2次元累積和の前計算をします。
 7        :math:`O(hw)` です。
 8
 9        Args:
10          h (int): 行数です。
11          w (int): 列数です。
12          a (list[list[int]]):  ``CumulativeSum2D`` を構築する配列です。
13        """
14        acc = [0] * ((h + 1) * (w + 1))
15        for ij in range(h * w):
16            i, j = divmod(ij, w)
17            acc[(i + 1) * (w + 1) + j + 1] = (
18                acc[i * (w + 1) + j + 1]
19                + acc[(i + 1) * (w + 1) + j]
20                - acc[i * (w + 1) + j]
21                + a[i][j]
22            )
23        self.h = h
24        self.w = w
25        self.acc = acc
26
27    def sum(self, h1: int, w1: int, h2: int, w2: int) -> int:
28        """長方形領域 ``[h1, h2) x [w1, w2)`` の総和を返します。
29        :math:`O(1)` です。
30        """
31        assert (
32            h1 <= h2
33        ), f"IndexError: {self.__class__.__name__}.sum({h1}, {w1}, {h2}, {w2}), h={self.h}"
34        assert (
35            w1 <= w2
36        ), f"IndexError: {self.__class__.__name__}.sum({h1}, {w1}, {h2}, {w2}), w={self.w}"
37        return (
38            self.acc[h2 * (self.w + 1) + w2]
39            - self.acc[h2 * (self.w + 1) + w1]
40            - self.acc[h1 * (self.w + 1) + w2]
41            + self.acc[h1 * (self.w + 1) + w1]
42        )

仕様

class CumulativeSum2D(h: int, w: int, a: list[list[int]])[source]

Bases: object

2次元累積和です。

sum(h1: int, w1: int, h2: int, w2: int) int[source]

長方形領域 [h1, h2) x [w1, w2) の総和を返します。 \(O(1)\) です。