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