#685824 (PyPy3) No.1718 Random Squirrel

• 問題: No.1718 Random Squirrel
• ユーザー: zkou
• 提出日時: 2021-08-05 17:48:49
• 言語: PyPy3 (7.3.15)
• 結果: AC
• 実行時間: 521 ms / 2,000 ms
• コード長: 2,566 bytes
• コンパイル時間: 219 ms
• コンパイル使用メモリ: 82,192 KB
• 実行使用メモリ: 120,740 KB
• 最終ジャッジ日時: 2024-09-23 03:47:35
• 合計ジャッジ時間: 8,044 ms
• ジャッジサーバーID （参考情報）: judge4 / judge3

ファイルパターン	結果
sample	AC * 2
other	AC * 31

ソースコード

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
INF = 1001001001
N, K = map(int, input().split())
adj = [[] for _ in range(N)]
for _ in range(N - 1):
    u, v = map(int, input().split())
    u -= 1; v -= 1
    adj[u].append(v)
    adj[v].append(u)
Ds = list(map(int, input().split()))
for i in range(K):
    Ds[i] -= 1
in_D = [False] * N
for D in Ds:
    in_D[D] = True
# 最小シュタイナー木みたいなのを作る。
reserved = [False] * N
root = Ds[0]
parent = [-1] * N
stack = [root]
while stack:
    v = stack.pop()
    if in_D[v]:
        reserved[v] = True
        while parent[v] != -1 and not reserved[parent[v]]:
            v = parent[v]
            reserved[v] = True
    for nv in adj[v]:
        if nv == parent[v]:
            continue
        parent[nv] = v
        stack.append(nv)
# print(reserved)
# 各頂点から、
# 最小シュタイナー木みたいなのにたどりつくための距離と、
# たどりつく先の頂点を求める。
dist_Steiner = [INF] * N
nearest = [-1] * N
for v in range(N):
    if reserved[v]:
        dist_Steiner[v] = 0
        nearest[v] = v
stack = []
for v in range(N):
    if reserved[v]:
        for nv in adj[v]:
            if not reserved[nv]:
                stack.append(v)
                break
while stack:
    v = stack.pop()
    for nv in adj[v]:
        if dist_Steiner[nv] > dist_Steiner[v] + 1:
            dist_Steiner[nv] = dist_Steiner[v] + 1
            nearest[nv] = nearest[v]
            stack.append(nv)
# シュタイナー木を一周するときに、帰りの分をさぼれる。
# そのさぼった分の最大化をする。
def dist_from_root_in_reserved(root):
    dist_diameter = [INF] * N
    dist_diameter[root] = 0
    stack = [root]
    while stack:
        v = stack.pop()
        for nv in adj[v]:
            if reserved[nv] and dist_diameter[nv] > dist_diameter[v] + 1:
                dist_diameter[nv] = dist_diameter[v] + 1
                stack.append(nv)
    for i in range(N):
        if dist_diameter[i] == INF:
            dist_diameter[i] = -1
    return dist_diameter
dist_0 = dist_from_root_in_reserved(Ds[0])
end1 = max(range(N), key = lambda i: dist_0[i])
dist_end1 = dist_from_root_in_reserved(end1)
end2 = max(range(N), key = lambda i: dist_end1[i])
dist_end2 = dist_from_root_in_reserved(end2)
max_dist = [max(e1, e2) for e1, e2 in zip(dist_end1, dist_end2)]
# print(max_dist)
# print(dist_0)
# print(dist_end1)
# print(dist_end2)
loop_len = (sum(reserved) - 1) * 2
for v in range(N):
    s = nearest[v]
    answer = dist_Steiner[v] + loop_len - max_dist[s]
    print(answer)
 
 
הההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההה
XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX