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