1# from titan_pylib.data_structures.segment_quadratic_division.segment_quadratic_division import SegmentQuadraticDivision
2from typing import Union, Callable, TypeVar, Generic, Iterable
3from functools import reduce
4from itertools import chain
5
6T = TypeVar("T")
7
8
9class SegmentQuadraticDivision(Generic[T]):
10
11 def __init__(self, n_or_a: Union[int, Iterable[T]], op: Callable[[T, T], T], e: T):
12 if isinstance(n_or_a, int):
13 self.n = n_or_a
14 a = [e] * self.n
15 else:
16 a = list(n_or_a)
17 self.n = len(a)
18 self.op = op
19 self.e = e
20 self.size = int(self.n**0.5) + 1
21 self.bucket_cnt = (self.n + self.size - 1) // self.size
22 self.data = [
23 a[k * self.size : (k + 1) * self.size] for k in range(self.bucket_cnt)
24 ]
25 self.bucket_data = [reduce(self.op, v) for v in self.data]
26
27 def prod(self, l: int, r: int) -> T:
28 """Return op([l, r)). / 0 <= l <= r <= n / O(√N)"""
29 assert 0 <= l <= r <= self.n
30 if l == r:
31 return self.e
32 k1 = l // self.size
33 k2 = r // self.size
34 l -= k1 * self.size
35 r -= k2 * self.size
36 if k1 == k2:
37 s = reduce(self.op, self.data[k1][l:r])
38 else:
39 s = self.e
40 if l < len(self.data[k1]):
41 s = reduce(self.op, self.data[k1][l:])
42 if k1 + 1 < k2:
43 s = (
44 reduce(self.op, self.bucket_data[k1 + 1 : k2])
45 if s == self.e
46 else reduce(self.op, self.bucket_data[k1 + 1 : k2], s)
47 )
48 if k2 < self.bucket_cnt and r > 0:
49 s = (
50 reduce(self.op, self.data[k2][:r])
51 if s == self.e
52 else reduce(self.op, self.data[k2][:r], s)
53 )
54 return s
55
56 def all_prod(self) -> T:
57 """Return op([0, n)). / O(√N)"""
58 return reduce(self.op, self.bucket_data)
59
60 def __getitem__(self, indx: int) -> T:
61 k = indx // self.size
62 return self.data[k][indx - k * self.size]
63
64 def __setitem__(self, indx: int, key: T):
65 k = indx // self.size
66 self.data[k][indx - k * self.size] = key
67 self.bucket_data[k] = reduce(self.op, self.data[k])
68
69 def __str__(self):
70 return str(list(chain(*self.data)))
71
72 def __repr__(self):
73 return f"SegmentQuadraticDivision({self})"