1from titan_pylib.data_structures.wbt._wbt_lazy_list_node import _WBTLazyListNode
2from typing import Generic, TypeVar, Optional, Iterable, Callable
3
4T = TypeVar("T")
5F = TypeVar("F")
6
7
[docs]
8class WBTLazyList(Generic[T, F]):
9
10 def __init__(
11 self,
12 op: Callable[[T, T], T],
13 mapping: Callable[[F, T], T],
14 composition: Callable[[F, F], F],
15 e: T,
16 id_: F,
17 a: Iterable[T] = [],
18 ) -> None:
19 self._root = None
20 self._op = op
21 self._mapping = mapping
22 self._composition = composition
23 self._e = e
24 self._id = id_
25 self.__build(a)
26
27 def __build(self, a: Iterable[T]) -> None:
28 def build(
29 l: int, r: int, pnode: Optional[_WBTLazyListNode] = None
30 ) -> _WBTLazyListNode:
31 if l == r:
32 return None
33 mid = (l + r) // 2
34 node = _WBTLazyListNode(self, a[mid], self._id)
35 node._left = build(l, mid, node)
36 node._right = build(mid + 1, r, node)
37 node._par = pnode
38 node._update()
39 return node
40
41 if not isinstance(a, list):
42 a = list(a)
43 if not a:
44 return
45 self._root = build(0, len(a))
46
47 @classmethod
48 def _weight(self, node: Optional[_WBTLazyListNode]) -> int:
49 return node._size + 1 if node else 1
50
51 def _merge_with_root(
52 self,
53 l: Optional[_WBTLazyListNode],
54 root: _WBTLazyListNode,
55 r: Optional[_WBTLazyListNode],
56 ) -> _WBTLazyListNode:
57 if self._weight(l) * _WBTLazyListNode.DELTA < self._weight(r):
58 r._propagate()
59 r._left = self._merge_with_root(l, root, r._left)
60 r._left._par = r
61 r._par = None
62 r._update()
63 if self._weight(r._right) * _WBTLazyListNode.DELTA < self._weight(r._left):
64 return r._balance_right()
65 return r
66 elif self._weight(r) * _WBTLazyListNode.DELTA < self._weight(l):
67 l._propagate()
68 l._right = self._merge_with_root(l._right, root, r)
69 l._right._par = l
70 l._par = None
71 l._update()
72 if self._weight(l._left) * _WBTLazyListNode.DELTA < self._weight(l._right):
73 return l._balance_left()
74 return l
75 else:
76 root._left = l
77 root._right = r
78 if l:
79 l._par = root
80 if r:
81 r._par = root
82 root._update()
83 return root
84
85 def _split_node(
86 self, node: _WBTLazyListNode, k: int
87 ) -> tuple[Optional[_WBTLazyListNode], Optional[_WBTLazyListNode]]:
88 if not node:
89 return None, None
90 node._propagate()
91 par = node._par
92 u = k if node._left is None else k - node._left._size
93 s, t = None, None
94 if u == 0:
95 s = node._left
96 t = self._merge_with_root(None, node, node._right)
97 elif u < 0:
98 s, t = self._split_node(node._left, k)
99 t = self._merge_with_root(t, node, node._right)
100 else:
101 s, t = self._split_node(node._right, u - 1)
102 s = self._merge_with_root(node._left, node, s)
103 if s:
104 s._par = par
105 if t:
106 t._par = par
107 return s, t
108
[docs]
109 def find_order(self, k: int) -> "WBTLazyList[T, F]":
110 if k < 0:
111 k += len(self)
112 node = self._root
113 while True:
114 node._propagate()
115 t = node._left._size if node._left else 0
116 if t == k:
117 return node
118 if t < k:
119 k -= t + 1
120 node = node._right
121 else:
122 node = node._left
123
[docs]
124 def split(self, k: int) -> tuple["WBTLazyList", "WBTLazyList"]:
125 lnode, rnode = self._split_node(self._root, k)
126 l, r = (
127 WBTLazyList(self._op, self._mapping, self._composition, self._e, self._id),
128 WBTLazyList(self._op, self._mapping, self._composition, self._e, self._id),
129 )
130 l._root = lnode
131 r._root = rnode
132 return l, r
133
134 def _pop_max(
135 self, node: _WBTLazyListNode
136 ) -> tuple[_WBTLazyListNode, _WBTLazyListNode]:
137 l, tmp = self._split_node(node, node._size - 1)
138 return l, tmp
139
140 def _merge_node(self, l: _WBTLazyListNode, r: _WBTLazyListNode) -> _WBTLazyListNode:
141 if l is None:
142 return r
143 if r is None:
144 return l
145 l, tmp = self._pop_max(l)
146 return self._merge_with_root(l, tmp, r)
147
[docs]
148 def extend(self, other: "WBTLazyList[T, F]") -> None:
149 self._root = self._merge_node(self._root, other._root)
150
[docs]
151 def insert(self, k: int, key) -> None:
152 s, t = self._split_node(self._root, k)
153 self._root = self._merge_with_root(s, _WBTLazyListNode(self, key, self._id), t)
154
[docs]
155 def pop(self, k: int):
156 s, t = self._split_node(self._root, k + 1)
157 s, tmp = self._pop_max(s)
158 self._root = self._merge_node(s, t)
159 return tmp._key
160
161 def _check(self, verbose: bool = False) -> None:
162 """作業用デバック関数
163 size,key,balanceをチェックして、正しければ高さを表示する
164 """
165 if self._root is None:
166 if verbose:
167 print("ok. 0 (empty)")
168 return
169
170 # _size, height
171 def dfs(node: _WBTLazyListNode) -> tuple[int, int]:
172 h = 0
173 s = 1
174 if node._left:
175 assert node._left._par is node
176 ls, lh = dfs(node._left)
177 s += ls
178 h = max(h, lh)
179 if node._right:
180 assert node._right._par is node
181 rs, rh = dfs(node._right)
182 s += rs
183 h = max(h, rh)
184 assert node._size == s
185 node._balance_check()
186 return s, h + 1
187
188 assert self._root._par is None
189 _, h = dfs(self._root)
190 if verbose:
191 print(f"ok. {h}")
192
[docs]
193 def prod(self, l: int, r: int) -> T:
194 def dfs(node: _WBTLazyListNode[T, F], left: int, right: int) -> T:
195 if right <= l or r <= left:
196 return self._e
197 node._propagate()
198 if l <= left and right < r:
199 return node._data
200 lsize = node._left._size if node._left else 0
201 res = self._e
202 if node._left:
203 res = dfs(node._left, left, left + lsize)
204 if l <= left + lsize < r:
205 res = self._op(res, node._key)
206 if node._right:
207 res = self._op(res, dfs(node._right, left + lsize + 1, right))
208 return res
209
210 return dfs(self._root, 0, len(self))
211
[docs]
212 def apply(self, l, r, f):
213 s, t = self._split_node(self._root, r)
214 r, s = self._split_node(s, l)
215 s._apply_lazy(f)
216 self._root = self._merge_node(self._merge_node(r, s), t)
217
[docs]
218 def reverse(self, l, r):
219 s, t = self._split_node(self._root, r)
220 r, s = self._split_node(s, l)
221 s._apply_rev()
222 self._root = self._merge_node(self._merge_node(r, s), t)
223
224 def __len__(self):
225 return self._root._size if self._root else 0
226
227 def __iter__(self):
228 node = self._root
229 stack: list[_WBTLazyListNode] = []
230 while stack or node:
231 if node:
232 node._propagate()
233 stack.append(node)
234 node = node._left
235 else:
236 node = stack.pop()
237 yield node._key
238 node = node._right
239
240 def __str__(self):
241 return str(list(self))