1# from titan_pylib.data_structures.b_tree.b_tree_list import BTreeList
2from collections import deque
3from typing import Generic, TypeVar, Iterable
4
5T = TypeVar("T")
6
7
8class BTreeList(Generic[T]):
9
10 class _Node:
11
12 def __init__(self):
13 self.key: list[T] = []
14 self.bit_len: list[int] = []
15 self.key_len: int = 0
16 self.child: list["BTreeList._Node"] = []
17 self.size: int = 0
18 self.sum: int = 0
19
20 def is_leaf(self) -> bool:
21 return not self.child
22
23 def _add_size(self, s: int) -> None:
24 self.size += s
25
26 def split(self, i: int) -> "BTreeList._Node":
27 right = BTreeList._Node()
28 self.key, right.key = self.key[:i], self.key[i:]
29 self.child, right.child = self.child[: i + 1], self.child[i + 1 :]
30 size = len(self.key) + sum(cnode.size for cnode in self.child)
31 s = sum(self.key) + sum(cnode.sum for cnode in self.child)
32 right.sum = self.sum - s
33 right.size = self.size - size
34 self.sum = s
35 self.size = size
36 return right
37
38 def insert_key(self, i: int, key: T, size: int = -1) -> None:
39 self.size += 1
40 self.sum += key if size == -1 else size
41 self.key.insert(i, key)
42
43 def insert_child(self, i: int, node: "BTreeList._Node", size=-1) -> None:
44 self.size += node.size if size == -1 else size
45 self.sum += node.sum if size == -1 else size
46 self.child.insert(i, node)
47
48 def append_key(self, key: T) -> None:
49 self.size += 1
50 self.sum += key
51 self.key.append(key)
52
53 def append_child(self, node: "BTreeList._Node", size: int = -1) -> None:
54 self.size += node.size if size == -1 else size
55 self.sum += node.sum if size == -1 else size
56 self.child.append(node)
57
58 def pop_key(self, i: int = -1) -> T:
59 self.size -= 1
60 v = self.key.pop(i)
61 self.sum -= v
62 return v
63
64 def len_key(self) -> int:
65 return len(self.key)
66
67 def pop_child(self, i: int = -1, size: int = -1) -> "BTreeList._Node":
68 cnode = self.child.pop(i)
69 self.sum -= cnode.sum if size == -1 else size
70 self.size -= cnode.size if size == -1 else size
71 return cnode
72
73 def extend_key(self, keys: list[T]) -> None:
74 self.size += len(keys)
75 self.sum += sum(keys)
76 self.key += keys
77
78 def extend_child(self, children: list["BTreeList._Node"]) -> None:
79 self.size += sum(cnode.size for cnode in children)
80 self.sum += sum(cnode.sum for cnode in children)
81 self.child += children
82
83 def __str__(self):
84 return str((str(self.key), self.size, self.sum))
85
86 __repr__ = __str__
87
88 def __init__(self, a: Iterable[T] = []):
89 self._m = 300
90 self._root = BTreeList._Node()
91 self._len = 0
92 self._build(a)
93
94 def _build(self, a: Iterable[T]):
95 for e in a:
96 self.append(e)
97
98 def __getitem__(self, k: int) -> T:
99 node = self._root
100 while True:
101 if node.is_leaf():
102 return node.key[k]
103 for i in range(node.len_key() + 1):
104 if k < node.child[i].size:
105 node = node.child[i]
106 break
107 k -= node.child[i].size
108 if k == 0 and i < node.len_key():
109 return node.key[i]
110 k -= 1
111
112 def pref(self, r: int) -> int:
113 node = self._root
114 s = 0
115 while True:
116 if node.is_leaf():
117 s += sum(node.key[:r])
118 break
119 for i in range(node.len_key() + 1):
120 if r < node.child[i].size:
121 node = node.child[i]
122 break
123 s += node.child[i].sum
124 r -= node.child[i].size
125 if r == 0 and i < node.len_key():
126 break
127 if i < node.len_key():
128 s += node.key[i]
129 r -= 1
130 return s
131
132 def prod(self, l: int, r: int) -> int:
133 return self.pref(r) - self.pref(l)
134
135 def _is_over(self, node: "BTreeList._Node") -> bool:
136 return node.len_key() > self._m
137
138 def insert(self, k: int, key: T) -> bool:
139 node = self._root
140 stack = []
141 while True:
142 if node.is_leaf():
143 node.insert_key(k, key)
144 break
145 for i in range(node.len_key() + 1):
146 if k < node.child[i].size:
147 break
148 k -= node.child[i].size
149 if k == 0:
150 k = node.child[i].size
151 break
152 k -= 1
153 stack.append((node, i))
154 node = node.child[i]
155 self._len += 1
156 while stack:
157 if not self._is_over(node):
158 break
159 pnode, indx = stack.pop()
160 i = node.len_key() // 2
161 pre = node.sum
162 center = node.pop_key(i)
163 right = node.split(i)
164 pnode.insert_key(indx, center, size=0)
165 pnode.insert_child(indx + 1, right, size=0)
166 pnode.sum += key
167 node = pnode
168 while stack:
169 pnode, _ = stack.pop()
170 pnode._add_size(1)
171 pnode.sum += key
172 if self._is_over(node):
173 pnode = BTreeList._Node()
174 i = node.len_key() // 2
175 center = node.pop_key(i)
176 right = node.split(i)
177 pnode.append_key(center)
178 pnode.append_child(node)
179 pnode.append_child(right)
180 self._root = pnode
181 return True
182
183 def append(self, key: T) -> None:
184 node = self._root
185 stack = []
186 while True:
187 if node.is_leaf():
188 node.append_key(key)
189 break
190 stack.append((node, node.len_key()))
191 node = node.child[-1]
192 self._len += 1
193 while stack:
194 if not self._is_over(node):
195 break
196 pnode, indx = stack.pop()
197 i = node.len_key() // 2
198 pre = node.sum
199 center = node.pop_key(i)
200 right = node.split(i)
201 pnode.insert_key(indx, center, size=0)
202 pnode.insert_child(indx + 1, right, size=0)
203 pnode.sum += key
204 node = pnode
205 while stack:
206 pnode, _ = stack.pop()
207 pnode._add_size(1)
208 pnode.sum += key
209 if self._is_over(node):
210 pnode = BTreeList._Node()
211 i = node.len_key() // 2
212 center = node.pop_key(i)
213 right = node.split(i)
214 pnode.append_key(center)
215 pnode.append_child(node)
216 pnode.append_child(right)
217 self._root = pnode
218 return True
219
220 def _discard_right(self, node: "BTreeList._Node") -> T:
221 stack = []
222 while not node.is_leaf():
223 stack.append(node)
224 if node.child[-1].len_key() == self._m // 2:
225 if node.child[-2].len_key() > self._m // 2:
226 cnode = node.child[-2]
227 node.child[-1].insert_key(0, node.key[-1])
228 node.key[-1] = cnode.pop_key()
229 if cnode.child:
230 node.child[-1].insert_child(0, cnode.pop_child())
231 node = node.child[-1]
232 continue
233 cnode = self._merge(node, node.len_key() - 1)
234 if node is self._root and not node.key:
235 self._root = cnode
236 node = cnode
237 continue
238 node = node.child[-1]
239 v = node.pop_key()
240 self._update_stack(stack, v)
241 return v
242
243 def _discard_left(self, node: "BTreeList._Node") -> T:
244 stack = []
245 while not node.is_leaf():
246 stack.append(node)
247 if node.child[0].len_key() == self._m // 2:
248 if node.child[1].len_key() > self._m // 2:
249 cnode = node.child[1]
250 node.child[0].append_key(node.key[0])
251 node.key[0] = cnode.pop_key(0)
252 if cnode.child:
253 node.child[0].append_child(cnode.pop_child(0))
254 node = node.child[0]
255 continue
256 cnode = self._merge(node, 0)
257 if node is self._root and not node.key:
258 self._root = cnode
259 node = cnode
260 continue
261 node = node.child[0]
262 v = node.pop_key(0)
263 self._update_stack(stack, v)
264 return v
265
266 def _merge(self, node: "BTreeList._Node", i: int) -> "BTreeList._Node":
267 y = node.child[i]
268 z = node.pop_child(i + 1, size=0)
269 v = node.pop_key(i)
270 y.append_key(v)
271 node.sum += v
272 node.size += 1
273 y.extend_key(z.key)
274 y.extend_child(z.child)
275 return y
276
277 def _merge_key(self, node: "BTreeList._Node", i: int) -> None:
278 if node.child[i].len_key() > self._m // 2:
279 node.key[i] = self._discard_right(node.child[i])
280 return
281 if node.child[i + 1].len_key() > self._m // 2:
282 node.key[i] = self._discard_left(node.child[i + 1])
283 return
284 indx = node.child[i if i + 1 < len(node.child) else (i - 1)].size
285 y = self._merge(node, i)
286 self._pop(indx, y)
287 if node is self._root and not node.key:
288 self._root = y
289
290 def _update_stack(self, stack, key):
291 for s in stack:
292 s.size -= 1
293 s.sum -= key
294
295 def tolist(self) -> list[T]:
296 a = []
297
298 def dfs(node):
299 if not node.child:
300 a.extend(node.key)
301 return
302 dfs(node.child[0])
303 for i in range(node.len_key()):
304 a.append(node.key[i])
305 dfs(node.child[i + 1])
306
307 dfs(self._root)
308 return a
309
310 def debug(self):
311 dep = [[] for _ in range(10)]
312 dq = deque([(self._root, 0)])
313 while dq:
314 node, d = dq.popleft()
315 dep[d].append((node.key, node.size, node.sum))
316 if node.child:
317 print(node, "child=", node.child)
318 for e in node.child:
319 if e:
320 dq.append((e, d + 1))
321 for i in range(10):
322 if not dep[i]:
323 break
324 for e in dep[i]:
325 print(e, end=" ")
326 print()
327
328 def pop(self, k: int) -> T:
329 assert 0 <= k < len(self), f"IndexError"
330 self._len -= 1
331 return self._pop(k, self._root)
332
333 def _pop(self, k: int, node: "BTreeList._Node") -> T:
334 stack = []
335 while True:
336 if node.is_leaf():
337 v = node.pop_key(k)
338 self._update_stack(stack, v)
339 return v
340 stack.append(node)
341 for i in range(node.len_key() + 1):
342 if k < node.child[i].size:
343 break
344 k -= node.child[i].size
345 if k == 0 and i < node.len_key():
346 v = node.key[i]
347 self._merge_key(node, i)
348 self._update_stack(stack, v)
349 return v
350 k -= 1
351 if node.child[i].len_key() == self._m // 2:
352 if (
353 i + 1 < len(node.child)
354 and node.child[i + 1].len_key() > self._m // 2
355 ):
356 cnode = node.child[i + 1]
357 node.child[i].append_key(node.key[i])
358 node.key[i] = cnode.pop_key(0)
359 if cnode.child:
360 node.child[i].append_child(cnode.pop_child(0))
361 node = node.child[i]
362 continue
363 if i - 1 >= 0 and node.child[i - 1].len_key() > self._m // 2:
364 cnode = node.child[i - 1]
365 k += 1
366 if cnode.child:
367 k += cnode.child[-1].size
368 node.child[i].insert_key(0, node.key[i - 1])
369 node.key[i - 1] = cnode.pop_key()
370 if cnode.child:
371 node.child[i].insert_child(0, cnode.pop_child())
372 node = node.child[i]
373 continue
374 if i + 1 >= len(node.child):
375 i -= 1
376 k += node.child[i].size + 1
377 cnode = self._merge(node, i)
378 if node is self._root and not node.key:
379 self._root = cnode
380 node = cnode
381 continue
382 node = node.child[i]
383
384 def __len__(self):
385 return self._len
386
387 def __str__(self):
388 return str(self.tolist())