1# from titan_pylib.data_structures.fenwick_tree.dynamic_fenwick_tree_2D import DynamicFenwickTree2D
2class DynamicFenwickTree2D:
3 """必要なところだけノードを作ります。2次元です。"""
4
5 def __init__(self, h: int, w: int, a: list[list[int]] = []):
6 """O(HWlogHlogW)"""
7 self._h: int = h + 1
8 self._w: int = w + 1
9 self._bit: dict[int, dict[int, int]] = {}
10 if a:
11 self._build(a)
12
13 def _build(self, a: list[list[int]]) -> None:
14 assert len(a) == self._h - 1 and len(a[0]) == self._w - 1
15 for i in range(self._h - 1):
16 for j in range(self._w - 1):
17 self.add(i, j, a[i][j])
18
19 def add(self, h: int, w: int, x) -> None:
20 """Add x to a[h][w]. / O(logH * logW)"""
21 h += 1
22 w += 1
23 _h, _w, _bit = self._h, self._w, self._bit
24 while h < _h:
25 j = w
26 if h not in _bit:
27 _bit[h] = {}
28 bit_h = _bit[h]
29 while j < _w:
30 if j in bit_h:
31 bit_h[j] += x
32 else:
33 bit_h[j] = x
34 j += j & -j
35 h += h & -h
36
37 def set(self, h: int, w: int, x) -> None:
38 self.add(h, w, x - self.get(h, w))
39
40 def _sum(self, h: int, w: int) -> int:
41 """Return sum([0, h) x [0, w)) of a. / O(logH * logW)"""
42 ret = 0
43 while h > 0:
44 j = w
45 if h not in self._bit:
46 h -= h & -h
47 continue
48 bit_h = self._bit[h]
49 while j > 0:
50 ret += bit_h.get(j, 0)
51 j -= j & -j
52 h -= h & -h
53 return ret
54
55 def sum(self, h1: int, w1: int, h2: int, w2: int) -> int:
56 """Retrun sum([h1, h2) x [w1, w2)) of a. / O(logH * logW)"""
57 assert h1 <= h2 and w1 <= w2
58 # w1, w2 = min(w1, w2), max(w1, w2)
59 # h1, h2 = min(h1, h2), max(h1, h2)
60 return (
61 self._sum(h2, w2)
62 - self._sum(h2, w1)
63 - self._sum(h1, w2)
64 + self._sum(h1, w1)
65 )
66
67 def get(self, h: int, w: int) -> int:
68 return self.sum(h, h + 1, w, w + 1)