Source code for titan_pylib.data_structures.fenwick_tree.fenwick_tree_2D_RAQ


[docs]
 1class FenwickTree2DRAQ:
 2    """2次元です。区間加算ができます"""
 3
 4    def __init__(self, h: int, w: int, a: list[list[int]] = []) -> None:
 5        """O(HWlogHlogW)"""
 6        self._h = h + 1
 7        self._w = w + 1
 8        self._bit = [[0] * (self._w) for _ in range(self._h)]
 9        if a:
10            assert len(a) == h
11            if h == 0:
12                return
13            assert len(a[0]) == w
14            self._build(a)
15
16    def _build(self, a: list[list[int]]) -> None:
17        for i in range(self._h - 1):
18            for j in range(self._w - 1):
19                if a[i][j] != 0:
20                    self.add(i, j, a[i][j])
21

[docs]
22    def add(self, h: int, w: int, x: int) -> None:
23        """Add x to a[h][w]. / O(logH * logW)"""
24        h += 1
25        w += 1
26        _h, _w, _bit = self._h, self._w, self._bit
27        while h < _h:
28            j = w
29            bit_h = _bit[h]
30            while j < _w:
31                bit_h[j] += x
32                j += j & -j
33            h += h & -h

34

[docs]
35    def set(self, h: int, w: int, x: int) -> None:
36        self.add(h, w, x - self.get(h, w))

37

[docs]
38    def range_add(self, h1: int, w1: int, h2: int, w2: int, x: int) -> int:
39        """Add x to [h1, h2) x [w1, w2) of a. / O(logH * logW)"""
40        assert h1 <= h2 and w1 <= w2
41        self.add(h1, w1, x)
42        self.add(h2, w1, -x)
43        self.add(h1, w2, -x)
44        self.add(h2, w2, x)

45

[docs]
46    def get(self, h: int, w: int) -> int:
47        """Return a[h][w]. / O(logH * logW)"""
48        ret = 0
49        while h > 0:
50            j = w
51            bit_h = self._bit[h]
52            while j > 0:
53                ret += bit_h[j]
54                j -= j & -j
55            h -= h & -h
56        return ret

57
58    def __str__(self) -> str:
59        ret = []
60        for i in range(self._h - 1):
61            ret.append(
62                ", ".join(map(str, ((self.get(i, j)) for j in range(self._w - 1))))
63            )
64        return "[\n  " + "\n  ".join(map(str, ret)) + "\n]"