Source code for titan_pylib.data_structures.sparse_table.sparse_table

 1from typing import Generic, TypeVar, Iterable, Callable
 2
 3T = TypeVar("T")
 4
 5

[docs]
 6class SparseTable(Generic[T]):
 7    """``SparseTable`` です。
 8
 9    静的であることに注意してください。
10    """
11
12    def __init__(self, a: Iterable[T], op: Callable[[T, T], T], e: T = None):
13        """
14        :math:`O(n\\log{n})` です。
15
16        Args:
17          a (Iterable[T]): ``SparseTable`` を構築する配列です。
18          op (Callable[[T, T], T]): 2項演算の関数です。
19          e (T, optional): 単位元です。
20        """
21        if not isinstance(a, list):
22            a = list(a)
23        self.size = len(a)
24        log = self.size.bit_length() - 1
25        data = [a] + [[]] * log
26        for i in range(log):
27            pre = data[i]
28            l = 1 << i
29            data[i + 1] = [op(pre[j], pre[j + l]) for j in range(len(pre) - l)]
30        self.data = data
31        self.op = op
32        self.e = e
33

[docs]
34    def prod(self, l: int, r: int) -> T:
35        """区間 ``[l, r)`` の総積を返します。
36
37        :math:`O(\\log{n})` です。
38
39        Args:
40          l (int): インデックスです。
41          r (int): インデックスです。
42        """
43        assert (
44            0 <= l <= r <= self.size
45        ), f"IndexError: {self.__class__.__name__}.prod({l}, {r}), len={self.size}"
46        if l == r:
47            return self.e
48        u = (r - l).bit_length() - 1
49        return self.op(self.data[u][l], self.data[u][r - (1 << u)])

50
51    def __getitem__(self, k: int) -> T:
52        assert (
53            0 <= k < self.size
54        ), f"IndexError: {self.__class__.__name__}[{k}], len={self.size}"
55        return self.data[0][k]
56
57    def __len__(self):
58        return self.size
59
60    def __str__(self):
61        return str(self.data[0])
62
63    def __repr__(self):
64        return f"{self.__class__.__name__}({self}, {self.op}, {self.e})"