cumulative_sum2D¶
ソースコード¶
from titan_pylib.data_structures.cumulative_sum.cumulative_sum2D import CumulativeSum2D
展開済みコード¶
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 )