1# from titan_pylib.data_structures.fenwick_tree.fenwick_tree_2D_RAQ import FenwickTree2DRAQ
2class FenwickTree2DRAQ:
3 """2次元です。区間加算ができます"""
4
5 def __init__(self, h: int, w: int, a: list[list[int]] = []) -> None:
6 """O(HWlogHlogW)"""
7 self._h = h + 1
8 self._w = w + 1
9 self._bit = [[0] * (self._w) for _ in range(self._h)]
10 if a:
11 assert len(a) == h
12 if h == 0:
13 return
14 assert len(a[0]) == w
15 self._build(a)
16
17 def _build(self, a: list[list[int]]) -> None:
18 for i in range(self._h - 1):
19 for j in range(self._w - 1):
20 if a[i][j] != 0:
21 self.add(i, j, a[i][j])
22
23 def add(self, h: int, w: int, x: int) -> None:
24 """Add x to a[h][w]. / O(logH * logW)"""
25 h += 1
26 w += 1
27 _h, _w, _bit = self._h, self._w, self._bit
28 while h < _h:
29 j = w
30 bit_h = _bit[h]
31 while j < _w:
32 bit_h[j] += x
33 j += j & -j
34 h += h & -h
35
36 def set(self, h: int, w: int, x: int) -> None:
37 self.add(h, w, x - self.get(h, w))
38
39 def range_add(self, h1: int, w1: int, h2: int, w2: int, x: int) -> int:
40 """Add x to [h1, h2) x [w1, w2) of a. / O(logH * logW)"""
41 assert h1 <= h2 and w1 <= w2
42 self.add(h1, w1, x)
43 self.add(h2, w1, -x)
44 self.add(h1, w2, -x)
45 self.add(h2, w2, x)
46
47 def get(self, h: int, w: int) -> int:
48 """Return a[h][w]. / O(logH * logW)"""
49 ret = 0
50 while h > 0:
51 j = w
52 bit_h = self._bit[h]
53 while j > 0:
54 ret += bit_h[j]
55 j -= j & -j
56 h -= h & -h
57 return ret
58
59 def __str__(self) -> str:
60 ret = []
61 for i in range(self._h - 1):
62 ret.append(
63 ", ".join(map(str, ((self.get(i, j)) for j in range(self._w - 1))))
64 )
65 return "[\n " + "\n ".join(map(str, ret)) + "\n]"