1from titan_pylib.data_structures.fenwick_tree.fenwick_tree import FenwickTree
2from titan_pylib.data_structures.segment_tree.segment_tree_RmQ import SegmentTreeRmQ
3
4
[docs]
5class EulerTour:
6
7 def __init__(
8 self, G: list[list[tuple[int, int]]], root: int, vertexcost: list[int] = []
9 ) -> None:
10 n = len(G)
11 if not vertexcost:
12 vertexcost = [0] * n
13
14 path = [0] * (2 * n)
15 vcost1 = [0] * (2 * n) # for vertex subtree
16 vcost2 = [0] * (2 * n) # for vertex path
17 ecost1 = [0] * (2 * n) # for edge subtree
18 ecost2 = [0] * (2 * n) # for edge path
19 nodein = [0] * n
20 nodeout = [0] * n
21 depth = [-1] * n
22
23 curtime = -1
24 depth[root] = 0
25 stack: list[tuple[int, int]] = [(~root, 0), (root, 0)]
26 while stack:
27 curtime += 1
28 v, ec = stack.pop()
29 if v >= 0:
30 nodein[v] = curtime
31 path[curtime] = v
32 ecost1[curtime] = ec
33 ecost2[curtime] = ec
34 vcost1[curtime] = vertexcost[v]
35 vcost2[curtime] = vertexcost[v]
36 if len(G[v]) == 1:
37 nodeout[v] = curtime + 1
38 for x, c in G[v]:
39 if depth[x] != -1:
40 continue
41 depth[x] = depth[v] + 1
42 stack.append((~v, c))
43 stack.append((x, c))
44 else:
45 v = ~v
46 path[curtime] = v
47 ecost1[curtime] = 0
48 ecost2[curtime] = -ec
49 vcost1[curtime] = 0
50 vcost2[curtime] = -vertexcost[v]
51 nodeout[v] = curtime
52
53 # ---------------------- #
54
55 self._n = n
56 self._depth = depth
57 self._nodein = nodein
58 self._nodeout = nodeout
59 self._vertexcost = vertexcost
60 self._path = path
61
62 self._vcost_subtree = FenwickTree(vcost1)
63 self._vcost_path = FenwickTree(vcost2)
64 self._ecost_subtree = FenwickTree(ecost1)
65 self._ecost_path = FenwickTree(ecost2)
66
67 bit = len(path).bit_length()
68 self.msk = (1 << bit) - 1
69 a: list[int] = [(depth[v] << bit) + i for i, v in enumerate(path)]
70 self._st: SegmentTreeRmQ[int] = SegmentTreeRmQ(a, e=max(a))
71
[docs]
72 def lca(self, u: int, v: int) -> int:
73 if u == v:
74 return u
75 l = min(self._nodein[u], self._nodein[v])
76 r = max(self._nodeout[u], self._nodeout[v])
77 ind = self._st.prod(l, r) & self.msk
78 return self._path[ind]
79
[docs]
80 def lca_mul(self, a: list[int]) -> int:
81 l, r = self._n + 1, -self._n - 1
82 for e in a:
83 l = min(l, self._nodein[e])
84 r = max(r, self._nodeout[e])
85 ind = self._st.prod(l, r) & self.msk
86 return self._path[ind]
87
[docs]
88 def subtree_vcost(self, v: int) -> int:
89 l = self._nodein[v]
90 r = self._nodeout[v]
91 return self._vcost_subtree.prod(l, r)
92
[docs]
93 def subtree_ecost(self, v: int) -> int:
94 l = self._nodein[v]
95 r = self._nodeout[v]
96 return self._ecost_subtree.prod(l + 1, r)
97
98 def _path_vcost(self, v: int) -> int:
99 """頂点 v を含む"""
100 return self._vcost_path.pref(self._nodein[v] + 1)
101
102 def _path_ecost(self, v: int) -> int:
103 """根から頂点 v までの辺"""
104 return self._ecost_path.pref(self._nodein[v] + 1)
105
[docs]
106 def path_vcost(self, u: int, v: int) -> int:
107 a = self.lca(u, v)
108 return (
109 self._path_vcost(u)
110 + self._path_vcost(v)
111 - 2 * self._path_vcost(a)
112 + self._vertexcost[a]
113 )
114
[docs]
115 def path_ecost(self, u: int, v: int) -> int:
116 return (
117 self._path_ecost(u)
118 + self._path_ecost(v)
119 - 2 * self._path_ecost(self.lca(u, v))
120 )
121
[docs]
122 def add_vertex(self, v: int, w: int) -> None:
123 """Add w to vertex x. / O(logN)"""
124 l = self._nodein[v]
125 r = self._nodeout[v]
126 self._vcost_subtree.add(l, w)
127 self._vcost_path.add(l, w)
128 self._vcost_path.add(r, -w)
129 self._vertexcost[v] += w
130
[docs]
131 def set_vertex(self, v: int, w: int) -> None:
132 """Set w to vertex v. / O(logN)"""
133 self.add_vertex(v, w - self._vertexcost[v])
134
[docs]
135 def add_edge(self, u: int, v: int, w: int) -> None:
136 """Add w to edge([u - v]). / O(logN)"""
137 if self._depth[u] < self._depth[v]:
138 u, v = v, u
139 l = self._nodein[u]
140 r = self._nodeout[u]
141 self._ecost_subtree.add(l, w)
142 self._ecost_subtree.add(r + 1, -w)
143 self._ecost_path.add(l, w)
144 self._ecost_path.add(r + 1, -w)
145
[docs]
146 def set_edge(self, u: int, v: int, w: int) -> None:
147 """Set w to edge([u - v]). / O(logN)"""
148 self.add_edge(u, v, w - self.path_ecost(u, v))