wbt_lazy_list

ソースコード

from titan_pylib.data_structures.wbt.wbt_lazy_list import WBTLazyList

view on github

展開済みコード

  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]

apply(l, r, f)[source]
extend(other: WBTLazyList[T, F]) None[source]
find_order(k: int) WBTLazyList[T, F][source]
insert(k: int, key) None[source]
pop(k: int)[source]
prod(l: int, r: int) T[source]
reverse(l, r)[source]
split(k: int) tuple[WBTLazyList, WBTLazyList][source]