1# from titan_pylib.data_structures.segment_tree.dual_segment_tree_commutative import DualSegmentTreeCommutative
2from typing import Union, Callable, TypeVar, Generic, Iterable
3
4T = TypeVar("T")
5F = TypeVar("F")
6
7
8class DualSegmentTreeCommutative(Generic[T, F]):
9 """双対セグ木です。"""
10
11 def __init__(
12 self,
13 n_or_a: Union[int, Iterable[T]],
14 mapping: Callable[[F, T], T],
15 composition: Callable[[F, F], F],
16 e: T,
17 id: F,
18 ) -> None:
19 self.mapping: Callable[[F, T], T] = mapping
20 self.composition: Callable[[F, F], F] = composition
21 self.e: T = e
22 self.id: F = id
23 self.data: list[T] = [e] * n_or_a if isinstance(n_or_a, int) else list(n_or_a)
24 self.n: int = len(self.data)
25 self.log: int = (self.n - 1).bit_length()
26 self.size: int = 1 << self.log
27 self.lazy: list[F] = [id] * self.size
28
29 def _all_apply(self, k: int, f: F) -> None:
30 if k < self.size:
31 self.lazy[k] = self.composition(f, self.lazy[k])
32 return
33 k -= self.size
34 if k < self.n:
35 self.data[k] = self.mapping(f, self.data[k])
36
37 def apply_point(self, k: int, f: F) -> None:
38 assert (
39 0 <= k < self.n
40 ), f"IndexError: {self.__class__.__name__}.apply_point({k}, {f}), n={self.n}"
41 k += self.size
42 self.data[k - self.size] = self.mapping(f, self.data[k - self.size])
43
44 def _propagate(self, k: int) -> None:
45 self._all_apply(k << 1, self.lazy[k])
46 self._all_apply(k << 1 | 1, self.lazy[k])
47 self.lazy[k] = self.id
48
49 def apply(self, l: int, r: int, f: F) -> None:
50 assert (
51 0 <= l <= r <= self.n
52 ), f"IndexError: {self.__class__.__name__}.apply({l}, {r}, {f}), n={self.n}"
53 if l == r:
54 return
55 if f == self.id:
56 return
57 l += self.size
58 r += self.size
59 lazy = self.lazy
60 l >>= 1
61 r >>= 1
62 while l < r:
63 if l & 1:
64 lazy[l] = self.composition(f, lazy[l])
65 l += 1
66 if r & 1:
67 r ^= 1
68 lazy[r] = self.composition(f, lazy[r])
69 l >>= 1
70 r >>= 1
71
72 def all_apply(self, f: F) -> None:
73 self.lazy[1] = self.composition(f, self.lazy[1])
74
75 def all_propagate(self) -> None:
76 for i in range(self.size):
77 self._propagate(i)
78
79 def tolist(self) -> list[T]:
80 self.all_propagate()
81 return self.data[:]
82
83 def __getitem__(self, k: int) -> T:
84 assert (
85 -self.n <= k < self.n
86 ), f"IndexError: {self.__class__.__name__}[{k}], n={self.n}"
87 if k < 0:
88 k += self.n
89 fs = self.id
90 k += self.size
91 for i in range(self.log, 0, -1):
92 fs = self.composition(fs, self.lazy[k >> i])
93 return self.mapping(fs, self.data[k - self.size])
94
95 def __setitem__(self, k: int, v: T) -> None:
96 assert (
97 -self.n <= k < self.n
98 ), f"IndexError: {self.__class__.__name__}[{k}] = {v}, n={self.n}"
99 if k < 0:
100 k += self.n
101 k += self.size
102 for i in range(self.log, 0, -1):
103 self._propagate(k >> i)
104 self.data[k - self.size] = v
105
106 def __str__(self):
107 return str([self[i] for i in range(self.n)])
108
109 def __repr__(self):
110 return f"{self.__class__.__name__}({self})"