1# from titan_pylib.data_structures.cumulative_sum.cumulative_op import CumulativeOp
2from typing import Generic, TypeVar, Callable, Iterable
3
4T = TypeVar("T")
5
6
7class CumulativeOp(Generic[T]):
8 """抽象化累積和です。"""
9
10 def __init__(
11 self, a: Iterable[T], op: Callable[[T, T], T], inv: Callable[[T], T], e: T
12 ):
13 """
14 :math:`O(n)` です。
15
16 Args:
17 a (Iterable[T]): ``CumulativeOp`` を構築する配列です。
18 op (Callable[[T, T], T]): 2項演算をする関数です。
19 inv (Callable[[T], T]): 逆元を返す関数です。 ``prod`` メソッドを使用する場合必要です。
20 e (T): 単位元です。
21 """
22 a = list(a)
23 n = len(a)
24 acc = [e for _ in range(n + 1)]
25 for i in range(n):
26 acc[i + 1] = op(acc[i], a[i])
27 self.n = n
28 self.acc = acc
29 self.a = a
30 self.op = op
31 self.inv = inv
32
33 def pref(self, r: int) -> T:
34 """区間 ``[0, r)`` の演算結果を返します。
35 :math:`O(1)` です。
36
37 Args:
38 r (int): インデックスです。
39 """
40 return self.acc[r]
41
42 def prod(self, l: int, r: int) -> T:
43 """区間 `[l, r)` の演算結果を返します。
44 :math:`O(1)` です。
45
46 Args:
47 l (int): インデックスです。
48 r (int): インデックスです。
49 """
50 assert (
51 0 <= l <= r <= self.n
52 ), f"IndexError: {self.__class__.__name__}.prod({l}, {r})"
53 return self.op(self.acc[r], self.inv(self.acc[l]))
54
55 def all_prod(self) -> T:
56 """区間 `[0, N)` の演算結果を返します。
57 :math:`O(1)` です。
58 """
59 return self.acc[-1]
60
61 def __getitem__(self, k: int) -> T:
62 return self.a[k]
63
64 def __len__(self):
65 return len(self.a)
66
67 def __str__(self):
68 return str(self.acc)
69
70 __repr__ = __str__