1from typing import Generic, TypeVar, Callable, Iterable, Optional
2
3T = TypeVar("T")
4F = TypeVar("F")
5
6
[docs]
7class LazyRBST(Generic[T, F]):
8
[docs]
9 class Random:
10
11 x: int = 2463534242
12
[docs]
13 @classmethod
14 def random32(cls) -> int:
15 cls.x ^= cls.x << 13 & 0xFFFFFFFF
16 cls.x ^= cls.x >> 17
17 cls.x ^= cls.x << 5 & 0xFFFFFFFF
18 return cls.x & 0xFFFFFFFF
19
[docs]
20 class Node:
21
22 def __init__(self, key: T, id: F):
23 self.key: T = key
24 self.data: T = key
25 self.lazy: F = id
26 self.left: Optional[LazyRBST.Node] = None
27 self.right: Optional[LazyRBST.Node] = None
28 self.size: int = 1
29 self.rev: int = 0
30
31 def __str__(self):
32 if self.left is None and self.right is None:
33 return f"key={self.key}, size={self.size}, data={self.data}, lazy={self.lazy}, rev={self.rev}\n"
34 return f"key={self.key}, size={self.size}, data={self.data}, lazy={self.lazy}, rev={self.rev},\n left:{self.left},\n right:{self.right}\n"
35
36 __repr__ = __str__
37
38 def __init__(
39 self,
40 a: Iterable[T],
41 op: Callable[[T, T], T],
42 mapping: Callable[[F, T], T],
43 composition: Callable[[F, F], F],
44 e: T,
45 id: F,
46 _root: Optional[Node] = None,
47 ):
48 self.root = _root
49 self.op = op
50 self.mapping = mapping
51 self.composition = composition
52 self.e = e
53 self.id = id
54 if a:
55 self.build(a)
56
[docs]
57 def build(self, a: Iterable[T]) -> None:
58 Node = LazyRBST.Node
59
60 def rec(l: int, r: int):
61 mid = (l + r) >> 1
62 node = Node(a[mid], id)
63 if l != mid:
64 node.left = rec(l, mid)
65 if mid + 1 != r:
66 node.right = rec(mid + 1, r)
67 self._update(node)
68 return node
69
70 a = list(a)
71 if not a:
72 return
73 id = self.id
74 self.root = rec(0, len(a))
75
76 def _propagate(self, node: Node) -> None:
77 if node.rev:
78 node.left, node.right = node.right, node.left
79 if node.left:
80 node.left.rev ^= 1
81 if node.right:
82 node.right.rev ^= 1
83 node.rev = 0
84 if node.lazy != self.id:
85 lazy = node.lazy
86 if node.left:
87 node.left.key = self.mapping(lazy, node.left.key)
88 node.left.data = self.mapping(lazy, node.left.data)
89 node.left.lazy = (
90 lazy
91 if node.left.lazy == self.id
92 else self.composition(lazy, node.left.lazy)
93 )
94 if node.right:
95 node.right.key = self.mapping(lazy, node.right.key)
96 node.right.data = self.mapping(lazy, node.right.data)
97 node.right.lazy = (
98 lazy
99 if node.right.lazy == self.id
100 else self.composition(lazy, node.right.lazy)
101 )
102 node.lazy = self.id
103
104 def _update(self, node: Node) -> None:
105 node.size = 1
106 node.data = node.key
107 if node.left:
108 node.size += node.left.size
109 node.data = self.op(node.left.data, node.key)
110 if node.right:
111 node.size += node.right.size
112 node.data = self.op(node.data, node.right.data)
113
114 def _update_lr(self, l: Node, r: Node):
115 l.size += r.size
116 l.data = self.op(l.data, r.data)
117
118 def _merge_node(self, l: Optional[Node], r: Optional[Node]) -> Optional[Node]:
119 root = None
120 r_root = None
121 d = -1
122 rand = LazyRBST.Random.random32
123 while l and r:
124 nd = rand() % (l.size + r.size) < l.size
125 node = l if nd else r
126 self._propagate(node)
127 if not root:
128 r_root = node
129 elif d == 0:
130 root.left = node
131 else:
132 root.right = node
133 root = node
134 d = nd
135 if d:
136 self._update_lr(l, r)
137 l = l.right
138 else:
139 self._update_lr(r, l)
140 r = r.left
141 if not root:
142 return l if l else r
143 if d == 0:
144 root.left = l if l else r
145 else:
146 root.right = l if l else r
147 return r_root
148
[docs]
149 def merge(self, other: "LazyRBST") -> None:
150 self.root = self._merge_node(self.root, other.root)
151
152 def _split_node(
153 self, node: Optional[Node], k: int
154 ) -> tuple[Optional[Node], Optional[Node]]:
155 left_path, right_path = [], []
156 while node:
157 self._propagate(node)
158 s = k if node.left is None else k - node.left.size
159 if s <= 0:
160 right_path.append(node)
161 node = node.left
162 else:
163 k = s - 1
164 left_path.append(node)
165 node = node.right
166 l, r = None, None
167 while left_path:
168 node = left_path.pop()
169 l, node.right = node, l
170 self._update(l)
171 while right_path:
172 node = right_path.pop()
173 r, node.left = node, r
174 self._update(r)
175 return l, r
176
[docs]
177 def split(self, k: int) -> tuple["LazyRBST", "LazyRBST"]:
178 left, right = self._split_node(self.root, k)
179 return (
180 LazyRBST(
181 [], self.op, self.mapping, self.composition, self.e, self.id, _root=left
182 ),
183 LazyRBST(
184 [],
185 self.op,
186 self.mapping,
187 self.composition,
188 self.e,
189 self.id,
190 _root=right,
191 ),
192 )
193
[docs]
194 def apply(self, l: int, r: int, f) -> None:
195 if l >= r:
196 return
197 s, t = self._split_node(self.root, r)
198 u, s = self._split_node(s, l)
199 assert s
200 s.key = self.mapping(f, s.key)
201 s.data = self.mapping(f, s.data)
202 s.lazy = self.composition(f, s.lazy)
203 self.root = self._merge_node(self._merge_node(u, s), t)
204
[docs]
205 def prod(self, l: int, r: int):
206 if l >= r:
207 return self.e
208 s, t = self._split_node(self.root, r)
209 u, s = self._split_node(s, l)
210 assert s
211 res = s.data
212 self.root = self._merge_node(self._merge_node(u, s), t)
213 return res
214
[docs]
215 def insert(self, k: int, key: T) -> None:
216 s, t = self._split_node(self.root, k)
217 self.root = self._merge_node(
218 self._merge_node(s, LazyRBST.Node(key, self.id)), t
219 )
220
[docs]
221 def pop(self, k: int) -> T:
222 s, t = self._split_node(self.root, k + 1)
223 assert s
224 s, tmp = self._split_node(s, s.size - 1)
225 res = tmp.key
226 self.root = self._merge_node(s, t)
227 return res
228
[docs]
229 def reverse(self, l: int, r: int) -> None:
230 if l >= r:
231 return
232 s, t = self._split_node(self.root, r)
233 u, s = self._split_node(s, l)
234 assert s
235 s.rev ^= 1
236 self.root = self._merge_node(self._merge_node(u, s), t)
237
[docs]
238 def tolist(self) -> list[T]:
239 node = self.root
240 stack, res = [], []
241 while stack or node:
242 if node:
243 self._propagate(node)
244 stack.append(node)
245 node = node.left
246 else:
247 node = stack.pop()
248 res.append(node.key)
249 node = node.right
250 return res
251
252 def __getitem__(self, k: int) -> T:
253 if k < 0:
254 k += len(self)
255 node = self.root
256 while True:
257 self._propagate(node)
258 t = 0 if node.left is None else node.left.size
259 if t == k:
260 return node.key
261 elif t < k:
262 k -= t + 1
263 node = node.right
264 else:
265 node = node.left
266
267 def __setitem__(self, k: int, key: T):
268 if k < 0:
269 k += len(self)
270 node = self.root
271 path = []
272 while True:
273 self._propagate(node)
274 path.append(node)
275 t = 0 if node.left is None else node.left.size
276 if t == k:
277 node.key = key
278 break
279 elif t < k:
280 k -= t + 1
281 node = node.right
282 else:
283 node = node.left
284 while path:
285 self._update(path.pop())
286
287 def __len__(self):
288 return 0 if self.root is None else self.root.size
289
290 def __str__(self):
291 return str(self.tolist())
292
293 __repr__ = __str__