1from titan_pylib.data_structures.segment_tree.segment_tree import SegmentTree
2from titan_pylib.graph.hld.hld import HLD
3from typing import Union, Iterable, Callable, TypeVar, Generic
4
5T = TypeVar("T")
6
7
[docs]
8class HLDNoncommutativeSegmentTree(Generic[T]):
9 """セグ木搭載HLDです。
10
11 Note:
12 **非可換に対応してます。**
13 """
14
15 def __init__(
16 self, hld: HLD, n_or_a: Union[int, Iterable[T]], op: Callable[[T, T], T], e: T
17 ):
18 self.hld: HLD = hld
19 a = (
20 [e] * n_or_a
21 if isinstance(n_or_a, int)
22 else self.hld.build_list(list(n_or_a))
23 )
24 self.seg: SegmentTree[T] = SegmentTree(a, op, e)
25 self.rseg: SegmentTree[T] = SegmentTree(a[::-1], op, e)
26 self.op: Callable[[T, T], T] = op
27 self.e: T = e
28
[docs]
29 def path_prod(self, u: int, v: int) -> T:
30 """頂点 ``u`` から頂点 ``v`` へのパスの集約値を返します。
31 :math:`O(\\log^2{n})` です。
32
33 Args:
34 u (int): パスの **始点** です。
35 v (int): パスの **終点** です。
36
37 Returns:
38 T: 求める集約値です。
39 """
40 head, nodein, dep, par, n = (
41 self.hld.head,
42 self.hld.nodein,
43 self.hld.dep,
44 self.hld.par,
45 self.hld.n,
46 )
47 lres, rres = self.e, self.e
48 seg, rseg = self.seg, self.rseg
49 while head[u] != head[v]:
50 if dep[head[u]] > dep[head[v]]:
51 lres = self.op(lres, rseg.prod(n - nodein[u] - 1, n - nodein[head[u]]))
52 u = par[head[u]]
53 else:
54 rres = self.op(seg.prod(nodein[head[v]], nodein[v] + 1), rres)
55 v = par[head[v]]
56 if dep[u] > dep[v]:
57 lres = self.op(lres, rseg.prod(n - nodein[u] - 1, n - nodein[v]))
58 else:
59 lres = self.op(lres, seg.prod(nodein[u], nodein[v] + 1))
60 return self.op(lres, rres)
61
[docs]
62 def get(self, k: int) -> T:
63 """頂点の値を返します。
64 :math:`O(\\log{n})` です。
65
66 Args:
67 k (int): 頂点のインデックスです。
68
69 Returns:
70 T: 頂点の値です。
71 """
72 return self.seg[self.hld.nodein[k]]
73
[docs]
74 def set(self, k: int, v: T) -> None:
75 """頂点の値を更新します。
76 :math:`O(\\log{n})` です。
77
78 Args:
79 k (int): 頂点のインデックスです。
80 v (T): 更新する値です。
81 """
82 self.seg[self.hld.nodein[k]] = v
83 self.rseg[self.hld.n - self.hld.nodein[k] - 1] = v
84
85 __getitem__ = get
86 __setitem__ = set
87
[docs]
88 def subtree_prod(self, v: int) -> T:
89 return self.seg.prod(self.hld.nodein[v], self.hld.nodeout[v])