1# from titan_pylib.data_structures.segment_tree.lazy_segment_tree_range import LazySegmentTreeRange
2from typing import Union, Callable, TypeVar, Generic, Iterable
3
4T = TypeVar("T")
5F = TypeVar("F")
6
7
8class LazySegmentTreeRange(Generic[T, F]):
9 """遅延セグ木です。"""
10
11 def __init__(
12 self,
13 n_or_a: Union[int, Iterable[T]],
14 op: Callable[[T, T, int], T],
15 mapping: Callable[[F, T, int], T],
16 composition: Callable[[F, F, int], F],
17 e: T,
18 id: F,
19 ) -> None:
20 self.op: Callable[[T, T, int], T] = op
21 self.mapping: Callable[[F, T, int], T] = mapping
22 self.composition: Callable[[F, F, int], F] = composition
23 self.e: T = e
24 self.id: F = id
25 if isinstance(n_or_a, int):
26 self.n = n_or_a
27 self.log = (self.n - 1).bit_length()
28 self.size = 1 << self.log
29 size_data: list[int] = [1] * (self.size << 1)
30 for i in range(self.size - 1, 0, -1):
31 size_data[i] = size_data[i << 1] + size_data[i << 1 | 1]
32 self.size_data = size_data
33 self.data = [e] * (self.size << 1)
34 else:
35 a = list(n_or_a)
36 self.n = len(a)
37 self.log = (self.n - 1).bit_length()
38 self.size = 1 << self.log
39 data = [e] * (self.size << 1)
40 data[self.size : self.size + self.n] = a
41 size_data: list[int] = [1] * (self.size << 1)
42 for i in range(self.size - 1, 0, -1):
43 size_data[i] = size_data[i << 1] + size_data[i << 1 | 1]
44 data[i] = op(data[i << 1], data[i << 1 | 1], size_data[i])
45 self.size_data = size_data
46 self.data = data
47 self.lazy = [id] * self.size
48
49 def _update(self, k: int) -> None:
50 self.data[k] = self.op(
51 self.data[k << 1], self.data[k << 1 | 1], self.size_data[k]
52 )
53
54 def _all_apply(self, k: int, f: F) -> None:
55 self.data[k] = self.mapping(f, self.data[k], self.size_data[k])
56 if k >= self.size:
57 return
58 self.lazy[k] = self.composition(f, self.lazy[k], self.size_data[k])
59
60 def _propagate(self, k: int) -> None:
61 self._all_apply(k << 1, self.lazy[k])
62 self._all_apply(k << 1 | 1, self.lazy[k])
63 self.lazy[k] = self.id
64
65 def apply_point(self, k: int, f: F) -> None:
66 k += self.size
67 for i in range(self.log, 0, -1):
68 self._propagate(k >> i)
69 self.data[k] = self.mapping(f, self.data[k], self.size_data[k])
70 for i in range(1, self.log + 1):
71 self._update(k >> i)
72
73 def apply(self, l: int, r: int, f: F) -> None:
74 assert (
75 0 <= l <= r <= self.n
76 ), f"IndexError: {self.__class__.__name__}.apply({l}, {r}, {f}), n={self.n}"
77 if l == r:
78 return
79 if f == self.id:
80 return
81 l += self.size
82 r += self.size
83 lazy = self.lazy
84 for i in range(self.log, 0, -1):
85 if l >> i << i != l and lazy[l >> i] != self.id:
86 self._propagate(l >> i)
87 if r >> i << i != r and lazy[(r - 1) >> i] != self.id:
88 self._propagate((r - 1) >> i)
89 l2, r2 = l, r
90 while l < r:
91 if l & 1:
92 self._all_apply(l, f)
93 l += 1
94 if r & 1:
95 self._all_apply(r ^ 1, f)
96 l >>= 1
97 r >>= 1
98 ll, rr = l2, r2 - 1
99 for i in range(1, self.log + 1):
100 ll >>= 1
101 rr >>= 1
102 if ll << i != l2:
103 self._update(ll)
104 if r2 >> i << i != r2:
105 self._update(rr)
106
107 def all_apply(self, f: F) -> None:
108 self.lazy[1] = self.composition(f, self.lazy[1], self.size_data[1])
109
110 def prod(self, l: int, r: int) -> T:
111 assert (
112 0 <= l <= r <= self.n
113 ), f"IndexError: {self.__class__.__name__}.prod({l}, {r}), n={self.n}"
114 if l == r:
115 return self.e
116 l += self.size
117 r += self.size
118 lazy = self.lazy
119 for i in range(self.log, 0, -1):
120 ll, rr = l >> i, r >> i
121 if ll << i != l and lazy[ll] != self.id:
122 self._propagate(ll)
123 if rr << i != r and lazy[rr] != self.id:
124 self._propagate(rr)
125 lres = self.e
126 rres = self.e
127 lsize = 1
128 rsize = 1
129 while l < r:
130 if l & 1:
131 lsize += self.size_data[l]
132 lres = self.op(lres, self.data[l], lsize)
133 l += 1
134 if r & 1:
135 rsize += self.size_data[r ^ 1]
136 rres = self.op(self.data[r ^ 1], rres, rsize)
137 l >>= 1
138 r >>= 1
139 return self.op(lres, rres, lsize + rsize)
140
141 def all_prod(self) -> T:
142 return self.data[1]
143
144 def all_propagate(self) -> None:
145 for i in range(self.size):
146 self._propagate(i)
147
148 def tolist(self) -> list[T]:
149 self.all_propagate()
150 return self.data[self.size : self.size + self.n]
151
152 def __getitem__(self, k: int) -> T:
153 assert (
154 -self.n <= k < self.n
155 ), f"IndexError: {self.__class__.__name__}[{k}], n={self.n}"
156 if k < 0:
157 k += self.n
158 k += self.size
159 for i in range(self.log, 0, -1):
160 self._propagate(k >> i)
161 return self.data[k]
162
163 def __setitem__(self, k: int, v: T):
164 assert (
165 -self.n <= k < self.n
166 ), f"IndexError: {self.__class__.__name__}[{k}] = {v}, n={self.n}"
167 if k < 0:
168 k += self.n
169 k += self.size
170 for i in range(self.log, 0, -1):
171 self._propagate(k >> i)
172 self.data[k] = v
173 for i in range(1, self.log + 1):
174 self._update(k >> i)
175
176 def __str__(self) -> str:
177 return "[" + ", ".join(map(str, (self[i] for i in range(self.n)))) + "]"
178
179 def __repr__(self):
180 return f"{self.__class__.__name__}({self})"