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