Source code for titan_pylib.data_structures.segment_quadratic_division.segment_quadratic_division

 1from typing import Union, Callable, TypeVar, Generic, Iterable
 2from functools import reduce
 3from itertools import chain
 4
 5T = TypeVar("T")
 6
 7

[docs]
 8class SegmentQuadraticDivision(Generic[T]):
 9
10    def __init__(self, n_or_a: Union[int, Iterable[T]], op: Callable[[T, T], T], e: T):
11        if isinstance(n_or_a, int):
12            self.n = n_or_a
13            a = [e] * self.n
14        else:
15            a = list(n_or_a)
16            self.n = len(a)
17        self.op = op
18        self.e = e
19        self.size = int(self.n**0.5) + 1
20        self.bucket_cnt = (self.n + self.size - 1) // self.size
21        self.data = [
22            a[k * self.size : (k + 1) * self.size] for k in range(self.bucket_cnt)
23        ]
24        self.bucket_data = [reduce(self.op, v) for v in self.data]
25

[docs]
26    def prod(self, l: int, r: int) -> T:
27        """Return op([l, r)). / 0 <= l <= r <= n / O(√N)"""
28        assert 0 <= l <= r <= self.n
29        if l == r:
30            return self.e
31        k1 = l // self.size
32        k2 = r // self.size
33        l -= k1 * self.size
34        r -= k2 * self.size
35        if k1 == k2:
36            s = reduce(self.op, self.data[k1][l:r])
37        else:
38            s = self.e
39            if l < len(self.data[k1]):
40                s = reduce(self.op, self.data[k1][l:])
41            if k1 + 1 < k2:
42                s = (
43                    reduce(self.op, self.bucket_data[k1 + 1 : k2])
44                    if s == self.e
45                    else reduce(self.op, self.bucket_data[k1 + 1 : k2], s)
46                )
47            if k2 < self.bucket_cnt and r > 0:
48                s = (
49                    reduce(self.op, self.data[k2][:r])
50                    if s == self.e
51                    else reduce(self.op, self.data[k2][:r], s)
52                )
53        return s

54

[docs]
55    def all_prod(self) -> T:
56        """Return op([0, n)). / O(√N)"""
57        return reduce(self.op, self.bucket_data)

58
59    def __getitem__(self, indx: int) -> T:
60        k = indx // self.size
61        return self.data[k][indx - k * self.size]
62
63    def __setitem__(self, indx: int, key: T):
64        k = indx // self.size
65        self.data[k][indx - k * self.size] = key
66        self.bucket_data[k] = reduce(self.op, self.data[k])
67
68    def __str__(self):
69        return str(list(chain(*self.data)))
70
71    def __repr__(self):
72        return f"SegmentQuadraticDivision({self})"