Source code for titan_pylib.data_structures.dynamic_connectivity.lazy_link_cut_tree

  1from array import array
  2from typing import Generic, TypeVar, Callable, Iterable, Union
  3from titan_pylib.data_structures.dynamic_connectivity.link_cut_tree import LinkCutTree
  4
  5T = TypeVar("T")
  6F = TypeVar("F")
  7
  8
[docs] 9class LazyLinkCutTree(LinkCutTree, Generic[T, F]): 10 """LazyLinkCutTree です。""" 11 12 # パスクエリ全部載せ 13 # - link / cut / merge / split 14 # - prod / apply / getitem / setitem 15 # - root / same 16 # - lca / path_length / path_kth_elm 17 # など 18 19 # opがいらないならupdateを即returnするように変更したり、 20 # 可換opならupdateを短縮したりなど 21 22 def __init__( 23 self, 24 n_or_a: Union[int, Iterable[T]], 25 op: Callable[[T, T], T], 26 mapping: Callable[[F, T], T], 27 composition: Callable[[F, F], F], 28 e: T, 29 id: F, 30 ) -> None: 31 """ 32 各引数は遅延セグ木のアレです。よしなに。 33 34 Args: 35 op (Callable[[T, T], T]): 非可換でも構いません。 36 """ 37 self.op = op 38 self.mapping = mapping 39 self.composition = composition 40 self.e = e 41 self.id = id 42 self.key: list[T] = [e] * (n_or_a) if isinstance(n_or_a, int) else list(n_or_a) 43 self.n = len(self.key) 44 self.key.append(e) 45 self.data: list[T] = [x for x in self.key for _ in range(2)] 46 self.lazy: list[F] = [id] * (self.n + 1) 47 self.arr: array[int] = array("I", [self.n, self.n, self.n, 0] * (self.n + 1)) 48 # node.left : arr[node<<2|0] 49 # node.right : arr[node<<2|1] 50 # node.par : arr[node<<2|2] 51 # node.rev : arr[node<<2|3] 52 self.size: array[int] = array("I", [1] * (self.n + 1)) 53 self.size[-1] = 0 54 self.group_cnt = self.n 55 56 def _propagate_lazy(self, node: int, f: F) -> None: 57 if node == self.n: 58 return 59 self.key[node] = self.mapping(f, self.key[node]) 60 self.data[node << 1] = self.mapping(f, self.data[node << 1]) 61 self.data[node << 1 | 1] = self.mapping(f, self.data[node << 1 | 1]) 62 self.lazy[node] = ( 63 f if self.lazy[node] == self.id else self.composition(f, self.lazy[node]) 64 ) 65 66 def _propagate(self, node: int) -> None: 67 if node == self.n: 68 return 69 arr = self.arr 70 if arr[node << 2 | 3]: 71 self.data[node << 1], self.data[node << 1 | 1] = ( 72 self.data[node << 1 | 1], 73 self.data[node << 1], 74 ) 75 arr[node << 2 | 3] = 0 76 ln, rn = arr[node << 2], arr[node << 2 | 1] 77 arr[node << 2] = rn 78 arr[node << 2 | 1] = ln 79 arr[ln << 2 | 3] ^= 1 80 arr[rn << 2 | 3] ^= 1 81 if self.lazy[node] == self.id: 82 return 83 self._propagate_lazy(self.arr[node << 2], self.lazy[node]) 84 self._propagate_lazy(self.arr[node << 2 | 1], self.lazy[node]) 85 self.lazy[node] = self.id 86 87 def _update(self, node: int) -> None: 88 if node == self.n: 89 return 90 ln, rn = self.arr[node << 2], self.arr[node << 2 | 1] 91 self._propagate(ln) 92 self._propagate(rn) 93 self.data[node << 1] = self.op( 94 self.op(self.data[ln << 1], self.key[node]), self.data[rn << 1] 95 ) 96 self.data[node << 1 | 1] = self.op( 97 self.op(self.data[rn << 1 | 1], self.key[node]), self.data[ln << 1 | 1] 98 ) 99 self.size[node] = 1 + self.size[ln] + self.size[rn] 100 101 def _update_triple(self, x: int, y: int, z: int) -> None: 102 data, key, arr, size = self.data, self.key, self.arr, self.size 103 lx, rx = arr[x << 2], arr[x << 2 | 1] 104 ly, ry = arr[y << 2], arr[y << 2 | 1] 105 self._propagate(lx) 106 self._propagate(rx) 107 self._propagate(ly) 108 self._propagate(ry) 109 data[z << 1] = data[x << 1] 110 data[x << 1] = self.op(self.op(data[lx << 1], key[x]), data[rx << 1]) 111 data[y << 1] = self.op(self.op(data[ly << 1], key[y]), data[ry << 1]) 112 data[z << 1 | 1] = data[x << 1 | 1] 113 data[x << 1 | 1] = self.op( 114 self.op(data[rx << 1 | 1], key[x]), data[lx << 1 | 1] 115 ) 116 data[y << 1 | 1] = self.op( 117 self.op(data[ry << 1 | 1], key[y]), data[ly << 1 | 1] 118 ) 119 size[z] = size[x] 120 size[x] = 1 + size[lx] + size[rx] 121 size[y] = 1 + size[ly] + size[ry] 122 123 def _update_double(self, x: int, y: int) -> None: 124 data, key, arr, size = self.data, self.key, self.arr, self.size 125 lx, rx = arr[x << 2], arr[x << 2 | 1] 126 self._propagate(lx) 127 self._propagate(rx) 128 data[y << 1] = data[x << 1] 129 data[x << 1] = self.op(self.op(data[lx << 1], key[x]), data[rx << 1]) 130 data[y << 1 | 1] = data[x << 1 | 1] 131 data[x << 1 | 1] = self.op( 132 self.op(data[rx << 1 | 1], key[x]), data[lx << 1 | 1] 133 ) 134 size[y] = size[x] 135 size[x] = 1 + size[lx] + size[rx] 136
[docs] 137 def path_prod(self, u: int, v: int) -> T: 138 """``u`` から ``v`` へのパスの総積を返します。 139 償却 :math:`O(\\log{n})` です。 140 """ 141 self.evert(u) 142 self.expose(v) 143 return self.data[v << 1]
144
[docs] 145 def path_apply(self, u: int, v: int, f: F) -> None: 146 """``u`` から ``v`` へのパスに ``f`` を作用させます。 147 償却 :math:`O(\\log{n})` です。 148 """ 149 self.evert(u) 150 self.expose(v) 151 self._propagate_lazy(v, f)
152
[docs] 153 def __setitem__(self, k: int, v: T): 154 """頂点 ``k`` の値を ``v`` に更新します。 155 償却 :math:`O(\\log{n})` です。 156 """ 157 self._splay(k) 158 self.key[k] = v 159 self._update(k)
160
[docs] 161 def __getitem__(self, k: int) -> T: 162 """頂点 ``k`` の値を返します。 163 償却 :math:`O(\\log{n})` です。 164 """ 165 self._splay(k) 166 return self.key[k]
167 168 def __str__(self): 169 return str([self[i] for i in range(self.n)]) 170 171 __repr__ = __str__