1from typing import Generic, Iterable, TypeVar, Callable, Optional
2
3T = TypeVar("T")
4F = TypeVar("F")
5
6
[docs]
7class ReversibleLazyAVLTree(Generic[T, F]):
8 """♰完全永続遅延伝播反転可能抽象化平衡二分探索木♰です。"""
9
[docs]
10 class Node:
11
12 def __init__(self, key: T, id: F):
13 self.key: T = key
14 self.data: T = key
15 self.rdata: T = key
16 self.left: Optional[ReversibleLazyAVLTree.Node] = None
17 self.right: Optional[ReversibleLazyAVLTree.Node] = None
18 self.lazy: F = id
19 self.rev: int = 0
20 self.height: int = 1
21 self.size: int = 1
22
23 def __str__(self):
24 if self.left is None and self.right is None:
25 return f"key:{self.key, self.height, self.size, self.data, self.rdata, self.lazy, self.rev}\n"
26 return f"key:{self.key, self.height, self.size, self.data, self.rdata, self.lazy, self.rev},\n left:{self.left},\n right:{self.right}\n"
27
28 def __init__(
29 self,
30 a: Iterable[T],
31 op: Callable[[T, T], T],
32 mapping: Callable[[F, T], T],
33 composition: Callable[[F, F], F],
34 e: T,
35 id: F,
36 node: Node = None,
37 ) -> None:
38 self.root: Optional[ReversibleLazyAVLTree.Node] = node
39 self.op: Callable[[T, T], T] = op
40 self.mapping: Callable[[F, T], T] = mapping
41 self.composition: Callable[[F, F], F] = composition
42 self.e: T = e
43 self.id: F = id
44 a = list(a)
45 if a:
46 self._build(a)
47
48 def _build(self, a: list[T]) -> None:
49 Node = ReversibleLazyAVLTree.Node
50 id = self.id
51
52 def sort(l: int, r: int) -> ReversibleLazyAVLTree.Node:
53 mid = (l + r) >> 1
54 node = Node(a[mid], id)
55 if l != mid:
56 node.left = sort(l, mid)
57 if mid + 1 != r:
58 node.right = sort(mid + 1, r)
59 self._update(node)
60 return node
61
62 self.root = sort(0, len(a))
63
64 def _propagate(self, node: Node) -> None:
65 l, r = node.left, node.right
66 if node.rev:
67 node.data, node.rdata = node.rdata, node.data
68 node.left, node.right = r, l
69 if l:
70 l.rev ^= 1
71 if r:
72 r.rev ^= 1
73 node.rev = 0
74 if node.lazy != self.id:
75 lazy = node.lazy
76 if l:
77 l.data = self.mapping(lazy, l.data)
78 l.rdata = self.mapping(lazy, l.rdata)
79 l.key = self.mapping(lazy, l.key)
80 l.lazy = lazy if l.lazy == self.id else self.composition(lazy, l.lazy)
81 if r:
82 r.data = self.mapping(lazy, r.data)
83 r.rdata = self.mapping(lazy, r.rdata)
84 r.key = self.mapping(lazy, r.key)
85 r.lazy = lazy if r.lazy == self.id else self.composition(lazy, r.lazy)
86 node.lazy = self.id
87
88 def _update(self, node: Node) -> None:
89 node.size = 1
90 node.data = node.key
91 node.rdata = node.key
92 node.height = 1
93 if node.left:
94 self._propagate(node.left)
95 node.size += node.left.size
96 node.data = self.op(node.left.data, node.data)
97 node.rdata = self.op(node.rdata, node.left.rdata)
98 node.height = max(node.left.height + 1, 1)
99 if node.right:
100 self._propagate(node.right)
101 node.size += node.right.size
102 node.data = self.op(node.data, node.right.data)
103 node.rdata = self.op(node.right.rdata, node.rdata)
104 node.height = max(node.height, node.right.height + 1)
105
106 def _get_balance(self, node: Node) -> int:
107 return (
108 (0 if node.right is None else -node.right.height)
109 if node.left is None
110 else (
111 node.left.height
112 if node.right is None
113 else node.left.height - node.right.height
114 )
115 )
116
117 def _balance_left(self, node: Node) -> Node:
118 self._propagate(node.left)
119 if node.left.left is None or node.left.left.height + 2 == node.left.height:
120 u = node.left.right
121 self._propagate(u)
122 node.left.right = u.left
123 u.left = node.left
124 node.left = u.right
125 u.right = node
126 self._update(u.left)
127 else:
128 u = node.left
129 node.left = u.right
130 u.right = node
131 self._update(u.right)
132 self._update(u)
133 return u
134
135 def _balance_right(self, node: Node) -> Node:
136 self._propagate(node.right)
137 if node.right.right is None or node.right.right.height + 2 == node.right.height:
138 u = node.right.left
139 self._propagate(u)
140 node.right.left = u.right
141 u.right = node.right
142 node.right = u.left
143 u.left = node
144 self._update(u.right)
145 else:
146 u = node.right
147 node.right = u.left
148 u.left = node
149 self._update(u.left)
150 self._update(u)
151 return u
152
153 def _kth_elm(self, k: int) -> T:
154 if k < 0:
155 k += len(self)
156 node = self.root
157 while True:
158 self._propagate(node)
159 t = 0 if node.left is None else node.left.size
160 if t == k:
161 return node.key
162 elif t < k:
163 k -= t + 1
164 node = node.right
165 else:
166 node = node.left
167
168 def _merge_with_root(self, l: Node, root: Node, r: Node) -> Node:
169 diff = (
170 (0 if r is None else -r.height)
171 if l is None
172 else (l.height if r is None else l.height - r.height)
173 )
174 if diff > 1:
175 self._propagate(l)
176 l.right = self._merge_with_root(l.right, root, r)
177 self._update(l)
178 if (
179 -l.right.height
180 if l.left is None
181 else l.left.height - l.right.height == -2
182 ):
183 return self._balance_right(l)
184 return l
185 elif diff < -1:
186 self._propagate(r)
187 r.left = self._merge_with_root(l, root, r.left)
188 self._update(r)
189 if (
190 r.left.height
191 if r.right is None
192 else r.left.height - r.right.height == 2
193 ):
194 return self._balance_left(r)
195 return r
196 else:
197 root.left = l
198 root.right = r
199 self._update(root)
200 return root
201
202 def _merge_node(self, l: Node, r: Node) -> Node:
203 if l is None:
204 return r
205 if r is None:
206 return l
207 l, tmp = self._pop_max(l)
208 return self._merge_with_root(l, tmp, r)
209
[docs]
210 def merge(self, other: "ReversibleLazyAVLTree") -> None:
211 self.root = self._merge_node(self.root, other.node)
212
213 def _pop_max(self, node: Node) -> tuple[Node, Node]:
214 self._propagate(node)
215 path = []
216 mx = node
217 while node.right:
218 path.append(node)
219 mx = node.right
220 node = node.right
221 self._propagate(node)
222 path.append(node.left)
223 for _ in range(len(path) - 1):
224 node = path.pop()
225 if node is None:
226 path[-1].right = None
227 self._update(path[-1])
228 continue
229 b = self._get_balance(node)
230 path[-1].right = (
231 self._balance_left(node)
232 if b == 2
233 else self._balance_right(node) if b == -2 else node
234 )
235 self._update(path[-1])
236 if path[0]:
237 b = self._get_balance(path[0])
238 path[0] = (
239 self._balance_left(path[0])
240 if b == 2
241 else self._balance_right(path[0]) if b == -2 else path[0]
242 )
243 mx.left = None
244 self._update(mx)
245 return path[0], mx
246
247 def _split_node(self, node: Node, k: int) -> tuple[Node, Node]:
248 if not node:
249 return None, None
250 self._propagate(node)
251 tmp = k if node.left is None else k - node.left.size
252 if tmp == 0:
253 return node.left, self._merge_with_root(None, node, node.right)
254 elif tmp < 0:
255 s, t = self._split_node(node.left, k)
256 return s, self._merge_with_root(t, node, node.right)
257 else:
258 s, t = self._split_node(node.right, tmp - 1)
259 return self._merge_with_root(node.left, node, s), t
260
[docs]
261 def split(self, k: int) -> tuple["ReversibleLazyAVLTree", "ReversibleLazyAVLTree"]:
262 l, r = self._split_node(self.root, k)
263 return ReversibleLazyAVLTree(
264 [], self.op, self.mapping, self.composition, self.e, self.id, l
265 ), ReversibleLazyAVLTree(
266 [], self.op, self.mapping, self.composition, self.e, self.id, r
267 )
268
[docs]
269 def insert(self, k: int, key: T) -> None:
270 s, t = self._split_node(self.root, k)
271 self.root = self._merge_with_root(
272 s, ReversibleLazyAVLTree.Node(key, self.id), t
273 )
274
[docs]
275 def pop(self, k: int) -> T:
276 s, t = self._split_node(self.root, k + 1)
277 s, tmp = self._pop_max(s)
278 self.root = self._merge_node(s, t)
279 return tmp.key
280
[docs]
281 def apply(self, l: int, r: int, f: F) -> None:
282 if l >= r or (not self.root):
283 return
284 stack = [(self.root), (self.root, 0, len(self))]
285 while stack:
286 if isinstance(stack[-1], tuple):
287 node, left, right = stack.pop()
288 if right <= l or r <= left:
289 continue
290 self._propagate(node)
291 if l <= left and right < r:
292 node.key = self.mapping(f, node.key)
293 node.data = self.mapping(f, node.data)
294 node.rdata = self.mapping(f, node.rdata)
295 node.lazy = (
296 f if node.lazy == self.id else self.composition(f, node.lazy)
297 )
298 else:
299 lsize = node.left.size if node.left else 0
300 stack.append(node)
301 if node.left:
302 stack.append((node.left, left, left + lsize))
303 if l <= left + lsize < r:
304 node.key = self.mapping(f, node.key)
305 if node.right:
306 stack.append((node.right, left + lsize + 1, right))
307 else:
308 self._update(stack.pop())
309
[docs]
310 def all_apply(self, f: F) -> None:
311 if not self.root:
312 return
313 self.root.key = self.mapping(f, self.root.key)
314 self.root.data = self.mapping(f, self.root.data)
315 self.root.rdata = self.mapping(f, self.root.rdata)
316 self.root.lazy = (
317 f if self.root.lazy == self.id else self.composition(f, self.root.lazy)
318 )
319
[docs]
320 def reverse(self, l: int, r: int) -> None:
321 if l >= r:
322 return
323 s, t = self._split_node(self.root, r)
324 r, s = self._split_node(s, l)
325 s.rev ^= 1
326 self._propagate(s)
327 self.root = self._merge_node(self._merge_node(r, s), t)
328
[docs]
329 def all_reverse(self) -> None:
330 if self.root is None:
331 return
332 self.root.rev ^= 1
333 self._propagate(self.root)
334
[docs]
335 def prod(self, l: int, r: int) -> T:
336 if l >= r or (not self.root):
337 return self.e
338
339 def dfs(node: ReversibleLazyAVLTree.Node, left: int, right: int) -> T:
340 if right <= l or r <= left:
341 return self.e
342 self._propagate(node)
343 if l <= left and right < r:
344 return node.data
345 lsize = node.left.size if node.left else 0
346 res = self.e
347 if node.left:
348 res = dfs(node.left, left, left + lsize)
349 if l <= left + lsize < r:
350 res = self.op(res, node.key)
351 if node.right:
352 res = self.op(res, dfs(node.right, left + lsize + 1, right))
353 return res
354
355 return dfs(self.root, 0, len(self))
356
[docs]
357 def all_prod(self) -> T:
358 return self.root.data if self.root else self.e
359
[docs]
360 def clear(self) -> None:
361 self.root = None
362
[docs]
363 def tolist(self) -> list[T]:
364 node = self.root
365 stack = []
366 a = []
367 while stack or node:
368 if node:
369 self._propagate(node)
370 stack.append(node)
371 node = node.left
372 else:
373 node = stack.pop()
374 a.append(node.key)
375 node = node.right
376 return a
377
378 def __len__(self):
379 return 0 if self.root is None else self.root.size
380
381 def __iter__(self):
382 self.__iter = 0
383 return self
384
385 def __next__(self):
386 if self.__iter == len(self):
387 raise StopIteration
388 res = self[self.__iter]
389 self.__iter += 1
390 return res
391
392 def __reversed__(self):
393 for i in range(len(self)):
394 yield self[-i - 1]
395
396 def __bool__(self):
397 return self.root is not None
398
399 def __getitem__(self, k: int) -> T:
400 return self._kth_elm(k)
401
402 def __setitem__(self, k, key: T):
403 if k < 0:
404 k += len(self)
405 node = self.root
406 path = []
407 while True:
408 self._propagate(node)
409 path.append(node)
410 t = 0 if node.left is None else node.left.size
411 if t == k:
412 node.key = key
413 break
414 if t < k:
415 k -= t + 1
416 node = node.right
417 else:
418 node = node.left
419 while path:
420 self._update(path.pop())
421
422 def __str__(self):
423 return "[" + ", ".join(map(str, self.tolist())) + "]"
424
425 def __repr__(self):
426 return f"ReversibleLazyAVLTree({self})"