木のパス上の値の集約(ダブリング)
LCA と同じく根からの BFS で親と深さを求め(再帰なし)、 $v$ の $2^k$ 個上の祖先と、そこまでの $2^k$ 本の辺の重みを op でまとめた値の表を作ります。 $u$ と $v$ を LCA まで持ち上げるときに飛んだ区間の値をまとめると、パス上の辺の重みをまとめた値になります。
op は結合的で可換な演算(最大値・最小値・和・gcd など)に限ります。パスの $u$ 側と $v$ 側を、向きや順番を気にせずにまとめるためです。
TreeDoubling(n, edges, op, e, root=0): 頂点数 $n$、木の辺のリストedges = [(u, v, w), ...]、演算opとその単位元e、根から前計算。 $O(N \log N)$query(u, v): $u$ と $v$ を結ぶパス上の辺の重みをopでまとめた値( $u = v$ ならe)。 $O(\log N)$lca(u, v): $u$ と $v$ の最小共通祖先。 $O(\log N)$up[k][v],val[k][v]: $v$ の $2^k$ 個上の祖先(根を越えるときは根)と、そこまでの辺の重みをopでまとめた値depth[v]: 根からの深さ(辺の本数)
op と e の例は、最大値なら max と十分小さい値、最小値なら min と十分大きい値、和なら lambda a, b: a + b と $0$、gcd なら math.gcd と $0$ です。
辺の重みが途中で変わるときは、HL分解 とセグメント木を組み合わせると、重みの更新とパス上の集約をどちらも $O(\log^2 N)$ で処理できます。
例(ABC 235 E。最小全域木に辺 $(u, v, w)$ を加えたとき、木の $u$ と $v$ を結ぶパス上の最大の重みが $w$ より大きければ、その辺を使う最小全域木がある。最小全域木(Kruskal・Prim) の kruskal を使う):
cost, used = kruskal(n, edges)
td = TreeDoubling(n, [edges[i] for i in used], max, 0) # 重みは正
for u, v, w in queries:
print("Yes" if w < td.query(u, v) else "No")
コード
class TreeDoubling:
def __init__(self, n, edges, op, e, root=0):
# edges = [(u, v, w), ...]。op は結合的かつ可換な演算、e はその単位元
g = [[] for _ in range(n)]
for u, v, w in edges:
g[u].append((v, w))
g[v].append((u, w))
self.op = op
self.e = e
self.depth = [-1] * n
self.depth[root] = 0
par = [root] * n # 根の親は根自身とする
pw = [e] * n # 親への辺の重み
order = [root]
for v in order:
for to, w in g[v]:
if self.depth[to] == -1:
self.depth[to] = self.depth[v] + 1
par[to] = v
pw[to] = w
order.append(to)
self.log = max(1, (n - 1).bit_length())
# up[k][v] = v の 2^k 個上の祖先、val[k][v] = そこまでの 2^k 本の辺の重みを op でまとめた値
self.up = [par]
self.val = [pw]
for _ in range(self.log - 1):
p, s = self.up[-1], self.val[-1]
self.up.append([p[p[v]] for v in range(n)])
self.val.append([op(s[v], s[p[v]]) for v in range(n)])
def _climb(self, u, v):
# (lca(u, v), u と v を結ぶパス上の辺の重みを op でまとめた値)
op, up, val = self.op, self.up, self.val
res = self.e
if self.depth[u] < self.depth[v]:
u, v = v, u
k = self.depth[u] - self.depth[v]
i = 0
while k:
if k & 1:
res = op(res, val[i][u])
u = up[i][u]
k >>= 1
i += 1
if u == v:
return u, res
for i in range(self.log - 1, -1, -1):
if up[i][u] != up[i][v]:
res = op(res, op(val[i][u], val[i][v]))
u = up[i][u]
v = up[i][v]
return up[0][u], op(res, op(val[0][u], val[0][v]))
def lca(self, u, v):
return self._climb(u, v)[0]
def query(self, u, v):
# u と v を結ぶパス上の辺の重みを op でまとめた値(u == v なら e)
return self._climb(u, v)[1]