1from titan_pylib.data_structures.segment_tree.lazy_segment_tree import LazySegmentTree
2from titan_pylib.graph.hld.hld import HLD
3from typing import Union, Iterable, Callable, TypeVar, Generic
4
5T = TypeVar("T")
6F = TypeVar("F")
7
8
[docs]
9class HLDLazySegmentTree(Generic[T, F]):
10 """遅延セグ木搭載HLDです。
11
12 Note:
13 **非可換に対応してます。**
14 """
15
16 def __init__(
17 self,
18 hld: HLD,
19 n_or_a: Union[int, Iterable[T]],
20 op: Callable[[T, T], T],
21 mapping: Callable[[F, T], T],
22 composition: Callable[[F, F], F],
23 e: T,
24 id: F,
25 ) -> None:
26 self.hld: HLD = hld
27 a = (
28 [e] * n_or_a
29 if isinstance(n_or_a, int)
30 else self.hld.build_list(list(n_or_a))
31 )
32 self.seg: LazySegmentTree[T, F] = LazySegmentTree(
33 a, op, mapping, composition, e, id
34 )
35 self.rseg: LazySegmentTree[T, F] = LazySegmentTree(
36 a[::-1], op, mapping, composition, e, id
37 )
38 self.op: Callable[[T, T], T] = op
39 self.e: T = e
40
[docs]
41 def path_prod(self, u: int, v: int) -> T:
42 """頂点 ``u`` から頂点 ``v`` へのパスの集約値を返します。
43 :math:`O(\\log^2{n})` です。
44
45 Args:
46 u (int): パスの **始点** です。
47 v (int): パスの **終点** です。
48
49 Returns:
50 T: 求める集約値です。
51 """
52 head, nodein, dep, par, n = (
53 self.hld.head,
54 self.hld.nodein,
55 self.hld.dep,
56 self.hld.par,
57 self.hld.n,
58 )
59 lres, rres = self.e, self.e
60 seg, rseg = self.seg, self.rseg
61 while head[u] != head[v]:
62 if dep[head[u]] > dep[head[v]]:
63 lres = self.op(lres, rseg.prod(n - nodein[u] - 1, n - nodein[head[u]]))
64 u = par[head[u]]
65 else:
66 rres = self.op(seg.prod(nodein[head[v]], nodein[v] + 1), rres)
67 v = par[head[v]]
68 if dep[u] > dep[v]:
69 lres = self.op(lres, rseg.prod(n - nodein[u] - 1, n - nodein[v]))
70 else:
71 lres = self.op(lres, seg.prod(nodein[u], nodein[v] + 1))
72 return self.op(lres, rres)
73
[docs]
74 def path_apply(self, u: int, v: int, f: F) -> None:
75 """頂点 ``u`` から頂点 ``v`` へのパスに作用させます。
76 :math:`O(\\log^2{n})` です。
77
78 Args:
79 u (int): パスの **始点** です。
80 v (int): パスの **終点** です。
81 f (F): 作用素です。
82 """
83 head, nodein, dep, par = (
84 self.hld.head,
85 self.hld.nodein,
86 self.hld.dep,
87 self.hld.par,
88 )
89 while head[u] != head[v]:
90 if dep[head[u]] < dep[head[v]]:
91 u, v = v, u
92 self.seg.apply(nodein[head[u]], nodein[u] + 1, f)
93 self.rseg.apply(
94 self.hld.n - (nodein[u] + 1 - 1) - 1,
95 self.hld.n - nodein[head[u]] - 1 + 1,
96 f,
97 )
98 u = par[head[u]]
99 if dep[u] < dep[v]:
100 u, v = v, u
101 self.seg.apply(nodein[v], nodein[u] + 1, f)
102 self.rseg.apply(
103 self.hld.n - (nodein[u] + 1 - 1) - 1, self.hld.n - nodein[v] - 1 + 1, f
104 )
105
[docs]
106 def get(self, k: int) -> T:
107 """頂点の値を返します。
108 :math:`O(\\log{n})` です。
109
110 Args:
111 k (int): 頂点のインデックスです。
112
113 Returns:
114 T: 頂点の値です。
115 """
116 return self.seg[self.hld.nodein[k]]
117
[docs]
118 def set(self, k: int, v: T) -> None:
119 """頂点の値を更新します。
120 :math:`O(\\log{n})` です。
121
122 Args:
123 k (int): 頂点のインデックスです。
124 v (T): 更新する値です。
125 """
126 self.seg[self.hld.nodein[k]] = v
127 self.rseg[self.hld.n - self.hld.nodein[k] - 1] = v
128
129 __getitem__ = get
130 __setitem__ = set
131
[docs]
132 def subtree_prod(self, v: int) -> T:
133 """部分木の集約値を返します。
134 :math:`O(\\log{n})` です。
135
136 Args:
137 v (int): 根とする頂点です。
138
139 Returns:
140 T: 求める集約値です。
141 """
142 return self.seg.prod(self.hld.nodein[v], self.hld.nodeout[v])
143
[docs]
144 def subtree_apply(self, v: int, f: F) -> None:
145 """部分木に作用させます。
146 :math:`O(\\log{n})` です。
147
148 Args:
149 v (int): 根とする頂点です。
150 f (F): 作用素です。
151 """
152 self.seg.apply(self.hld.nodein[v], self.hld.nodeout[v], f)
153 self.rseg.apply(
154 self.hld.n - self.hld.nodeout[v] - 1 - 1,
155 self.hld.n - self.hld.nodein[v] - 1 + 1,
156 f,
157 )