Source code for titan_pylib.graph.hld.hld_segment_tree

 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 HLDSegmentTree(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    ) -> None:
18        self.hld: HLD = hld
19        n_or_a = (
20            n_or_a if isinstance(n_or_a, int) else self.hld.build_list(list(n_or_a))
21        )
22        self.seg: SegmentTree[T] = SegmentTree(n_or_a, op, e)
23        self.op: Callable[[T, T], T] = op
24        self.e: T = e
25

[docs]
26    def path_prod(self, u: int, v: int) -> T:
27        """頂点 ``u`` から頂点 ``v`` へのパスの集約値を返します。
28        :math:`O(\\log^2{n})` です。
29
30        Note:
31            **非可換に対応していません。**
32
33        Args:
34            u (int): パスの端点です。
35            v (int): パスの端点です。
36
37        Returns:
38            T: 求める集約値です。
39        """
40        head, nodein, dep, par = (
41            self.hld.head,
42            self.hld.nodein,
43            self.hld.dep,
44            self.hld.par,
45        )
46        res = self.e
47        while head[u] != head[v]:
48            if dep[head[u]] < dep[head[v]]:
49                u, v = v, u
50            res = self.op(res, self.seg.prod(nodein[head[u]], nodein[u] + 1))
51            u = par[head[u]]
52        if dep[u] < dep[v]:
53            u, v = v, u
54        return self.op(res, self.seg.prod(nodein[v], nodein[u] + 1))

55

[docs]
56    def get(self, k: int) -> T:
57        """頂点の値を返します。
58        :math:`O(\\log{n})` です。
59
60        Args:
61            k (int): 頂点のインデックスです。
62
63        Returns:
64            T: 頂点の値です。
65        """
66        return self.seg[self.hld.nodein[k]]

67

[docs]
68    def set(self, k: int, v: T) -> None:
69        """頂点の値を更新します。
70        :math:`O(\\log{n})` です。
71
72        Args:
73            k (int): 頂点のインデックスです。
74            v (T): 更新する値です。
75        """
76        self.seg[self.hld.nodein[k]] = v

77
78    __getitem__ = get
79    __setitem__ = set
80

[docs]
81    def subtree_prod(self, v: int) -> T:
82        """部分木の集約値を返します。
83        :math:`O(\\log{n})` です。
84
85        Args:
86            v (int): 根とする頂点です。
87
88        Returns:
89            T: 求める集約値です。
90        """
91        return self.seg.prod(self.hld.nodein[v], self.hld.nodeout[v])