Source code for titan_pylib.graph.hld.hld_lazy_segment_tree

  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 )