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