Day 018-Q1 — 仮想木（Virtual Tree）

問題

N頂点の根付き木（根は頂点1）が与えられる。各辺には重みがある。Q個のクエリでK個の頂点集合 S が指定される。S の全頂点のLCAの関係を保ちつつ、S に含まれない頂点を圧縮した「仮想木」を構築し、仮想木上での辺重み総和を求めよ。

制約

入出力例

7
1 2 3
1 3 4
2 4 2
2 5 1
3 6 5
3 7 2
2
3 4 5 6
2 4 7

15
11

ヒント (段階的開示)

nodes = sorted(set(nodes), key=lambda x: euler_in[x])
extras = [lca(nodes[i], nodes[i+1]) for i in range(len(nodes)-1)]
all_nodes = sorted(set(nodes + extras), key=lambda x: euler_in[x])
# スタックで仮想木構築

模範解答 (Python)

📋 コピーimport sys
from collections import defaultdict
input = sys.stdin.readline

def main():
    sys.setrecursionlimit(300000)
    N = int(input())
    graph = defaultdict(list)
    for _ in range(N-1):
        u, v, w = map(int, input().split())
        u -= 1; v -= 1
        graph[u].append((v, w))
        graph[v].append((u, w))

    LOG = 18
    parent = [[-1]*N for _ in range(LOG)]
    depth = [0]*N
    dist = [0]*N
    euler_in = [0]*N
    euler_out = [0]*N
    timer = [0]

    stack = [(0, -1, 0, 0, False)]
    while stack:
        v, p, d, dd, returning = stack.pop()
        if returning:
            euler_out[v] = timer[0]; timer[0] += 1
            continue
        parent[0][v] = p if p != -1 else v
        depth[v] = d; dist[v] = dd
        euler_in[v] = timer[0]; timer[0] += 1
        stack.append((v, p, d, dd, True))
        for u, w in graph[v]:
            if u != p:
                stack.append((u, v, d+1, dd+w, False))

    for k in range(1, LOG):
        for v in range(N):
            parent[k][v] = parent[k-1][parent[k-1][v]]

    def lca(u, v):
        if depth[u] < depth[v]: u, v = v, u
        diff = depth[u] - depth[v]
        for k in range(LOG):
            if (diff >> k) & 1: u = parent[k][u]
        if u == v: return u
        for k in range(LOG-1, -1, -1):
            if parent[k][u] != parent[k][v]:
                u = parent[k][u]; v = parent[k][v]
        return parent[0][u]

    def edge_dist(u, v):
        l = lca(u, v)
        return dist[u] + dist[v] - 2*dist[l]

    Q = int(input())
    for _ in range(Q):
        line = list(map(int, input().split()))
        K = line[0]
        nodes = [x-1 for x in line[1:K+1]]
        nodes_sorted = sorted(set(nodes), key=lambda x: euler_in[x])
        to_add = []
        for i in range(len(nodes_sorted)-1):
            to_add.append(lca(nodes_sorted[i], nodes_sorted[i+1]))
        all_nodes = sorted(set(nodes_sorted + to_add), key=lambda x: euler_in[x])

        vtree_edges = []
        stack2 = []
        for v in all_nodes:
            if not stack2:
                stack2.append(v)
            else:
                l = lca(v, stack2[-1])
                if l != stack2[-1]:
                    while len(stack2) >= 2 and depth[stack2[-2]] >= depth[l]:
                        vtree_edges.append((stack2[-2], stack2[-1]))
                        stack2.pop()
                    if stack2[-1] != l:
                        vtree_edges.append((l, stack2[-1]))
                        stack2.pop()
                        stack2.append(l)
                stack2.append(v)
        while len(stack2) >= 2:
            vtree_edges.append((stack2[-2], stack2[-1]))
            stack2.pop()

        ans = sum(edge_dist(u, v) for u, v in vtree_edges)
        print(ans)

main()

Step-by-Step 解説

よくあるミス

次のステップ

• 発展: 仮想木上でのDPを実装（仮想木の辺を使った木DP）
• 応用: Steiner Tree問題への応用

自己評価

自分の回答

気づき・メモ