lazy_splay_tree_array¶
ソースコード¶
from titan_pylib.data_structures.splay_tree.lazy_splay_tree_array import LazySplayTreeArrayData
from titan_pylib.data_structures.splay_tree.lazy_splay_tree_array import LazySplayTreeArray
展開済みコード¶
1# from titan_pylib.data_structures.splay_tree.lazy_splay_tree_array import LazySplayTreeArray
2from array import array
3from typing import (
4 Generic,
5 TypeVar,
6 Callable,
7 Iterable,
8 Optional,
9 Union,
10 Sequence,
11)
12
13T = TypeVar("T")
14F = TypeVar("F")
15
16
17class LazySplayTreeArrayData(Generic[T, F]):
18
19 def __init__(
20 self,
21 op: Optional[Callable[[T, T], T]] = None,
22 mapping: Optional[Callable[[F, T], T]] = None,
23 composition: Optional[Callable[[F, F], F]] = None,
24 e: T = None,
25 id: F = None,
26 ):
27 self.op: Callable[[T, T], T] = (lambda s, t: e) if op is None else op
28 self.mapping: Callable[[F, T], T] = (lambda f, s: e) if op is None else mapping
29 self.composition: Callable[[F, F], F] = (
30 (lambda f, g: id) if op is None else composition
31 )
32 self.e: T = e
33 self.id: F = id
34 self.keydata: list[T] = [e, e]
35 self.lazy: list[F] = [id]
36 self.arr: array[int] = array("I", bytes(16))
37 # left: arr[node<<2]
38 # right: arr[node<<2|1]
39 # size: arr[node<<2|2]
40 # rev: arr[node<<2|3]
41 self.end: int = 1
42
43 def reserve(self, n: int) -> None:
44 if n <= 0:
45 return
46 self.keydata += [self.e] * (2 * n)
47 self.lazy += [self.id] * n
48 self.arr += array("I", bytes(16 * n))
49
50
51class LazySplayTreeArray(Generic[T, F]):
52
53 def __init__(
54 self,
55 data: "LazySplayTreeArrayData[T, F]",
56 n_or_a: Union[int, Iterable[T]] = 0,
57 _root: int = 0,
58 ) -> None:
59 self.data = data
60 self.root = _root
61 if not n_or_a:
62 return
63 if isinstance(n_or_a, int):
64 a = [data.e] * n_or_a
65 elif not isinstance(n_or_a, Sequence):
66 a = list(n_or_a)
67 else:
68 a = n_or_a
69 if a:
70 self._build(a)
71
72 def _build(self, a: Sequence[T]) -> None:
73 def rec(l: int, r: int) -> int:
74 mid = (l + r) >> 1
75 if l != mid:
76 arr[mid << 2] = rec(l, mid)
77 if mid + 1 != r:
78 arr[mid << 2 | 1] = rec(mid + 1, r)
79 self._update(mid)
80 return mid
81
82 n = len(a)
83 keydata, arr = self.data.keydata, self.data.arr
84 end = self.data.end
85 self.data.reserve(n + end - len(keydata) // 2 + 1)
86 self.data.end += n
87 for i, e in enumerate(a):
88 keydata[end + i << 1] = e
89 keydata[end + i << 1 | 1] = e
90 self.root = rec(end, n + end)
91
92 def _make_node(self, key: T) -> int:
93 data = self.data
94 if data.end >= len(data.arr) // 4:
95 data.keydata.append(key)
96 data.keydata.append(key)
97 data.lazy.append(data.id)
98 data.arr.append(0)
99 data.arr.append(0)
100 data.arr.append(1)
101 data.arr.append(0)
102 else:
103 data.keydata[data.end << 1] = key
104 data.keydata[data.end << 1 | 1] = key
105 data.end += 1
106 return data.end - 1
107
108 def _propagate(self, node: int) -> None:
109 data = self.data
110 arr = data.arr
111 if arr[node << 2 | 3]:
112 arr[node << 2], arr[node << 2 | 1] = arr[node << 2 | 1], arr[node << 2]
113 arr[node << 2 | 3] = 0
114 arr[arr[node << 2] << 2 | 3] ^= 1
115 arr[arr[node << 2 | 1] << 2 | 3] ^= 1
116 nlazy = data.lazy[node]
117 if nlazy == data.id:
118 return
119 lnode, rnode = arr[node << 2], arr[node << 2 | 1]
120 keydata, lazy = data.keydata, data.lazy
121 lazy[node] = data.id
122 if lnode:
123 lazy[lnode] = data.composition(nlazy, lazy[lnode])
124 lnode <<= 1
125 keydata[lnode] = data.mapping(nlazy, keydata[lnode])
126 keydata[lnode | 1] = data.mapping(nlazy, keydata[lnode | 1])
127 if rnode:
128 lazy[rnode] = data.composition(nlazy, lazy[rnode])
129 rnode <<= 1
130 keydata[rnode] = data.mapping(nlazy, keydata[rnode])
131 keydata[rnode | 1] = data.mapping(nlazy, keydata[rnode | 1])
132
133 def _update_triple(self, x: int, y: int, z: int) -> None:
134 data = self.data
135 keydata, arr = data.keydata, data.arr
136 lx, rx = arr[x << 2], arr[x << 2 | 1]
137 ly, ry = arr[y << 2], arr[y << 2 | 1]
138 arr[z << 2 | 2] = arr[x << 2 | 2]
139 arr[x << 2 | 2] = 1 + arr[lx << 2 | 2] + arr[rx << 2 | 2]
140 arr[y << 2 | 2] = 1 + arr[ly << 2 | 2] + arr[ry << 2 | 2]
141 keydata[z << 1 | 1] = keydata[x << 1 | 1]
142 keydata[x << 1 | 1] = data.op(
143 data.op(keydata[lx << 1 | 1], keydata[x << 1]), keydata[rx << 1 | 1]
144 )
145 keydata[y << 1 | 1] = data.op(
146 data.op(keydata[ly << 1 | 1], keydata[y << 1]), keydata[ry << 1 | 1]
147 )
148
149 def _update_double(self, x: int, y: int) -> None:
150 data = self.data
151 keydata, arr = data.keydata, data.arr
152 lx, rx = arr[x << 2], arr[x << 2 | 1]
153 arr[y << 2 | 2] = arr[x << 2 | 2]
154 arr[x << 2 | 2] = 1 + arr[lx << 2 | 2] + arr[rx << 2 | 2]
155 keydata[y << 1 | 1] = keydata[x << 1 | 1]
156 keydata[x << 1 | 1] = data.op(
157 data.op(keydata[lx << 1 | 1], keydata[x << 1]), keydata[rx << 1 | 1]
158 )
159
160 def _update(self, node: int) -> None:
161 data = self.data
162 keydata, arr = data.keydata, data.arr
163 lnode, rnode = arr[node << 2], arr[node << 2 | 1]
164 arr[node << 2 | 2] = 1 + arr[lnode << 2 | 2] + arr[rnode << 2 | 2]
165 keydata[node << 1 | 1] = data.op(
166 data.op(keydata[lnode << 1 | 1], keydata[node << 1]),
167 keydata[rnode << 1 | 1],
168 )
169
170 def _splay(self, path: list[int], d: int) -> None:
171 arr = self.data.arr
172 g = d & 1
173 while len(path) > 1:
174 pnode = path.pop()
175 gnode = path.pop()
176 f = d >> 1 & 1
177 node = arr[pnode << 2 | g ^ 1]
178 nnode = (pnode if g == f else node) << 2 | f
179 arr[pnode << 2 | g ^ 1] = arr[node << 2 | g]
180 arr[node << 2 | g] = pnode
181 arr[gnode << 2 | f ^ 1] = arr[nnode]
182 arr[nnode] = gnode
183 self._update_triple(gnode, pnode, node)
184 if not path:
185 return
186 d >>= 2
187 g = d & 1
188 arr[path[-1] << 2 | g ^ 1] = node
189 pnode = path.pop()
190 node = arr[pnode << 2 | g ^ 1]
191 arr[pnode << 2 | g ^ 1] = arr[node << 2 | g]
192 arr[node << 2 | g] = pnode
193 self._update_double(pnode, node)
194
195 def _kth_elm_splay(self, node: int, k: int) -> int:
196 arr = self.data.arr
197 if k < 0:
198 k += arr[node << 2 | 2]
199 d = 0
200 path = []
201 while True:
202 self._propagate(node)
203 t = arr[arr[node << 2] << 2 | 2]
204 if t == k:
205 if path:
206 self._splay(path, d)
207 return node
208 d = d << 1 | (t > k)
209 path.append(node)
210 node = arr[node << 2 | (t < k)]
211 if t < k:
212 k -= t + 1
213
214 def _left_splay(self, node: int) -> int:
215 if not node:
216 return 0
217 self._propagate(node)
218 arr = self.data.arr
219 if not arr[node << 2]:
220 return node
221 path = []
222 while arr[node << 2]:
223 path.append(node)
224 node = arr[node << 2]
225 self._propagate(node)
226 self._splay(path, (1 << len(path)) - 1)
227 return node
228
229 def _right_splay(self, node: int) -> int:
230 if not node:
231 return 0
232 self._propagate(node)
233 arr = self.data.arr
234 if not arr[node << 2 | 1]:
235 return node
236 path = []
237 while arr[node << 2 | 1]:
238 path.append(node)
239 node = arr[node << 2 | 1]
240 self._propagate(node)
241 self._splay(path, 0)
242 return node
243
244 def reserve(self, n: int) -> None:
245 self.data.reserve(n)
246
247 def merge(self, other: "LazySplayTreeArray[T, F]") -> None:
248 assert self.data is other.data
249 if not other.root:
250 return
251 if not self.root:
252 self.root = other.root
253 return
254 self.root = self._right_splay(self.root)
255 self.data.arr[self.root << 2 | 1] = other.root
256 self._update(self.root)
257
258 def split(
259 self, k: int
260 ) -> tuple["LazySplayTreeArray[T, F]", "LazySplayTreeArray[T, F]"]:
261 assert (
262 -len(self) < k <= len(self)
263 ), f"IndexError: LazySplayTreeArray.split({k}), len={len(self)}"
264 if k < 0:
265 k += len(self)
266 if k >= self.data.arr[self.root << 2 | 2]:
267 return self, LazySplayTreeArray(self.data, _root=0)
268 self.root = self._kth_elm_splay(self.root, k)
269 left = LazySplayTreeArray(self.data, _root=self.data.arr[self.root << 2])
270 self.data.arr[self.root << 2] = 0
271 self._update(self.root)
272 return left, self
273
274 def _internal_split(self, k: int) -> tuple[int, int]:
275 if k >= self.data.arr[self.root << 2 | 2]:
276 return self.root, 0
277 self.root = self._kth_elm_splay(self.root, k)
278 left = self.data.arr[self.root << 2]
279 self.data.arr[self.root << 2] = 0
280 self._update(self.root)
281 return left, self.root
282
283 def reverse(self, l: int, r: int) -> None:
284 assert (
285 0 <= l <= r <= len(self)
286 ), f"IndexError: LazySplayTreeArray.reverse({l}, {r}), len={len(self)}"
287 if l == r:
288 return
289 data = self.data
290 left, right = self._internal_split(r)
291 if l:
292 left = self._kth_elm_splay(left, l - 1)
293 data.arr[(data.arr[left << 2 | 1] if l else left) << 2 | 3] ^= 1
294 if right:
295 data.arr[right << 2] = left
296 self._update(right)
297 self.root = right if right else left
298
299 def all_reverse(self) -> None:
300 self.data.arr[self.root << 2 | 3] ^= 1
301
302 def apply(self, l: int, r: int, f: F) -> None:
303 assert (
304 0 <= l <= r <= len(self)
305 ), f"IndexError: LazySplayTreeArray.apply({l}, {r}), len={len(self)}"
306 data = self.data
307 left, right = self._internal_split(r)
308 keydata, lazy = data.keydata, data.lazy
309 if l:
310 left = self._kth_elm_splay(left, l - 1)
311 node = data.arr[left << 2 | 1] if l else left
312 keydata[node << 1] = data.mapping(f, keydata[node << 1])
313 keydata[node << 1 | 1] = data.mapping(f, keydata[node << 1 | 1])
314 lazy[node] = data.composition(f, lazy[node])
315 if l:
316 self._update(left)
317 if right:
318 data.arr[right << 2] = left
319 self._update(right)
320 self.root = right if right else left
321
322 def all_apply(self, f: F) -> None:
323 if not self.root:
324 return
325 data, node = self.data, self.root
326 data.keydata[node << 1] = data.mapping(f, data.keydata[node << 1])
327 data.keydata[node << 1 | 1] = data.mapping(f, data.keydata[node << 1 | 1])
328 data.lazy[node] = data.composition(f, data.lazy[node])
329
330 def prod(self, l: int, r: int) -> T:
331 assert (
332 0 <= l <= r <= len(self)
333 ), f"IndexError: LazySplayTreeArray.prod({l}, {r}), len={len(self)}"
334 data = self.data
335 left, right = self._internal_split(r)
336 if l:
337 left = self._kth_elm_splay(left, l - 1)
338 res = data.keydata[(data.arr[left << 2 | 1] if l else left) << 1 | 1]
339 if right:
340 data.arr[right << 2] = left
341 self._update(right)
342 self.root = right if right else left
343 return res
344
345 def all_prod(self) -> T:
346 return self.data.keydata[self.root << 1 | 1]
347
348 def insert(self, k: int, key: T) -> None:
349 assert (
350 -len(self) <= k <= len(self)
351 ), f"IndexError: LazySplayTreeArray.insert({k}, {key}), len={len(self)}"
352 if k < 0:
353 k += len(self)
354 data = self.data
355 node = self._make_node(key)
356 if not self.root:
357 self._update(node)
358 self.root = node
359 return
360 arr = data.arr
361 if k == data.arr[self.root << 2 | 2]:
362 arr[node << 2] = self._right_splay(self.root)
363 else:
364 node_ = self._kth_elm_splay(self.root, k)
365 if arr[node_ << 2]:
366 arr[node << 2] = arr[node_ << 2]
367 arr[node_ << 2] = 0
368 self._update(node_)
369 arr[node << 2 | 1] = node_
370 self._update(node)
371 self.root = node
372
373 def append(self, key: T) -> None:
374 data = self.data
375 node = self._right_splay(self.root)
376 self.root = self._make_node(key)
377 data.arr[self.root << 2] = node
378 self._update(self.root)
379
380 def appendleft(self, key: T) -> None:
381 node = self._left_splay(self.root)
382 self.root = self._make_node(key)
383 self.data.arr[self.root << 2 | 1] = node
384 self._update(self.root)
385
386 def pop(self, k: int = -1) -> T:
387 assert -len(self) <= k < len(self), f"IndexError: LazySplayTreeArray.pop({k})"
388 data = self.data
389 if k == -1:
390 node = self._right_splay(self.root)
391 self._propagate(node)
392 self.root = data.arr[node << 2]
393 return data.keydata[node << 1]
394 self.root = self._kth_elm_splay(self.root, k)
395 res = data.keydata[self.root << 1]
396 if not data.arr[self.root << 2]:
397 self.root = data.arr[self.root << 2 | 1]
398 elif not data.arr[self.root << 2 | 1]:
399 self.root = data.arr[self.root << 2]
400 else:
401 node = self._right_splay(data.arr[self.root << 2])
402 data.arr[node << 2 | 1] = data.arr[self.root << 2 | 1]
403 self.root = node
404 self._update(self.root)
405 return res
406
407 def popleft(self) -> T:
408 assert self, "IndexError: LazySplayTreeArray.popleft()"
409 node = self._left_splay(self.root)
410 self.root = self.data.arr[node << 2 | 1]
411 return self.data.keydata[node << 1]
412
413 def rotate(self, x: int) -> None:
414 # 「末尾をを削除し先頭に挿入」をx回
415 n = self.data.arr[self.root << 2 | 2]
416 l, self = self.split(n - (x % n))
417 self.merge(l)
418
419 def tolist(self) -> list[T]:
420 node = self.root
421 arr, keydata = self.data.arr, self.data.keydata
422 stack = []
423 res = []
424 while stack or node:
425 if node:
426 self._propagate(node)
427 stack.append(node)
428 node = arr[node << 2]
429 else:
430 node = stack.pop()
431 res.append(keydata[node << 1])
432 node = arr[node << 2 | 1]
433 return res
434
435 def clear(self) -> None:
436 self.root = 0
437
438 def __setitem__(self, k: int, key: T):
439 assert (
440 -len(self) <= k < len(self)
441 ), f"IndexError: LazySplayTreeArray.__setitem__({k})"
442 self.root = self._kth_elm_splay(self.root, k)
443 self.data.keydata[self.root << 1] = key
444 self._update(self.root)
445
446 def __getitem__(self, k: int) -> T:
447 assert (
448 -len(self) <= k < len(self)
449 ), f"IndexError: LazySplayTreeArray.__getitem__({k})"
450 self.root = self._kth_elm_splay(self.root, k)
451 return self.data.keydata[self.root << 1]
452
453 def __iter__(self):
454 self.__iter = 0
455 return self
456
457 def __next__(self):
458 if self.__iter == self.data.arr[self.root << 2 | 2]:
459 raise StopIteration
460 res = self[self.__iter]
461 self.__iter += 1
462 return res
463
464 def __reversed__(self):
465 for i in range(len(self)):
466 yield self[-i - 1]
467
468 def __len__(self):
469 return self.data.arr[self.root << 2 | 2]
470
471 def __str__(self):
472 return str(self.tolist())
473
474 def __bool__(self):
475 return self.root != 0
476
477 def __repr__(self):
478 return f"{self.__class__.__name__}({self})"
仕様¶
- class LazySplayTreeArray(data: LazySplayTreeArrayData[T, F], n_or_a: int | Iterable[T] = 0, _root: int = 0)[source]¶
Bases:
Generic[T,F]- merge(other: LazySplayTreeArray[T, F]) None[source]¶
- split(k: int) tuple[LazySplayTreeArray[T, F], LazySplayTreeArray[T, F]][source]¶