Source code for titan_pylib.data_structures.cumulative_sum.cumulative_op

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

[docs]
 6class CumulativeOp(Generic[T]):
 7    """抽象化累積和です。"""
 8
 9    def __init__(
10        self, a: Iterable[T], op: Callable[[T, T], T], inv: Callable[[T], T], e: T
11    ):
12        """
13        :math:`O(n)` です。
14
15        Args:
16          a (Iterable[T]): ``CumulativeOp`` を構築する配列です。
17          op (Callable[[T, T], T]): 2項演算をする関数です。
18          inv (Callable[[T], T]): 逆元を返す関数です。 ``prod`` メソッドを使用する場合必要です。
19          e (T): 単位元です。
20        """
21        a = list(a)
22        n = len(a)
23        acc = [e for _ in range(n + 1)]
24        for i in range(n):
25            acc[i + 1] = op(acc[i], a[i])
26        self.n = n
27        self.acc = acc
28        self.a = a
29        self.op = op
30        self.inv = inv
31

[docs]
32    def pref(self, r: int) -> T:
33        """区間 ``[0, r)`` の演算結果を返します。
34        :math:`O(1)` です。
35
36        Args:
37          r (int): インデックスです。
38        """
39        return self.acc[r]

40

[docs]
41    def prod(self, l: int, r: int) -> T:
42        """区間 `[l, r)` の演算結果を返します。
43        :math:`O(1)` です。
44
45        Args:
46          l (int): インデックスです。
47          r (int): インデックスです。
48        """
49        assert (
50            0 <= l <= r <= self.n
51        ), f"IndexError: {self.__class__.__name__}.prod({l}, {r})"
52        return self.op(self.acc[r], self.inv(self.acc[l]))

53

[docs]
54    def all_prod(self) -> T:
55        """区間 `[0, N)` の演算結果を返します。
56        :math:`O(1)` です。
57        """
58        return self.acc[-1]

59
60    def __getitem__(self, k: int) -> T:
61        return self.a[k]
62
63    def __len__(self):
64        return len(self.a)
65
66    def __str__(self):
67        return str(self.acc)
68
69    __repr__ = __str__