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__