結果

問題 No.1718 Random Squirrel
コンテスト
ユーザー LyricalMaestro
提出日時 2024-12-15 03:44:35
言語 PyPy3
(7.3.15)
結果
AC  
実行時間 686 ms / 2,000 ms
コード長 4,265 bytes
コンパイル時間 480 ms
コンパイル使用メモリ 82,420 KB
実行使用メモリ 126,452 KB
最終ジャッジ日時 2024-12-15 03:44:45
合計ジャッジ時間 9,536 ms
ジャッジサーバーID
（参考情報） 		judge3 / judge1
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 2
other AC * 31
権限があれば一括ダウンロードができます

ソースコード

diff # 

## https://yukicoder.me/problems/no/1718

from collections import deque

def bfs(N, next_nodes, s_list):
    dists = [-1] * N
    queue = deque()

    for s, d in s_list:
        dists[s] = d
        queue.append(s)
    while len(queue) > 0:
        v = queue.popleft()
        for w in next_nodes[v]:
            if dists[w] == -1:
                dists[w] = dists[v] + 1
                queue.append(w)
    return dists

def solve(N, new_next_nodes, base_root):
    next_max_dists = [[0] * len(new_next_nodes[v]) for v in range(N)]
    max_dists = [0] * N
    stack = deque()
    stack.append((base_root, 0))
    parents = [-2] * N
    parents[base_root] = -1
    while len(stack) > 0:
        v, index = stack.pop()
        if index > 0:
            w = new_next_nodes[v][index - 1]
            max_dists[v] = max(max_dists[v], max_dists[w] + 1)
            next_max_dists[v][index - 1] = max_dists[w] + 1

        while index < len(new_next_nodes[v]):
            w = new_next_nodes[v][index]
            if w == parents[v]:
                index += 1
                continue

            parents[w] = v
            stack.append((v, index + 1))
            stack.append((w, 0))
            break
    
    queue = deque()
    queue.append((base_root, -1))
    while len(queue) > 0:
        v, d = queue.popleft()
        p = parents[v]

        for index in range(len(new_next_nodes[v])):
            w = new_next_nodes[v][index]
            if w == p:
                max_dists[v] = max(max_dists[v], d + 1)
                next_max_dists[v][index] = d + 1
        
        m_dists = next_max_dists[v]
        f_m_dists = [0] * len(next_max_dists[v])
        f_m = -1
        for i in range(1, len(next_max_dists[v])):
            f_m = max(f_m, m_dists[i - 1])
            f_m_dists[i] = f_m

        b_m_dists = [0] * len(next_max_dists[v])
        b_m = -1
        for i in reversed(range(len(next_max_dists[v]) - 1)):
            b_m = max(b_m, m_dists[i + 1])
            b_m_dists[i] = b_m
        
        for index in range(len(new_next_nodes[v])):
            w = new_next_nodes[v][index]
            if w == p:
                continue

            m = max(b_m_dists[index], f_m_dists[index])
            queue.append((w, m))
    return max_dists



def main():
    N, K = map(int, input().split())
    next_nodes = [[] for _ in range(N)]
    for _ in range(N - 1):
        u, v = map(int, input().split())
        next_nodes[u - 1].append(v - 1)
        next_nodes[v - 1].append(u - 1)
    
    D = list(map(lambda x : int(x) - 1, input().split()))

    if K == 1:
        # D[0]からの距離を計算していくだけ
        dists = bfs(N, next_nodes, [(D[0], 0)])
        for d in dists:
            print(d)
        return


    # 最小シュタイナー木構築
    d_map = [False] * N
    for d in D:
        d_map[d] = True
    stack = deque()
    stack.append((D[0], 0))
    parents = [-2] * N
    parents[D[0]] = -1
    needs = [False] * N
    needs[D[0]] = True
    while len(stack ) > 0:
        v, index = stack.pop()
        if index == 0:
            if d_map[v]:
                needs[v] = True
        else:
            if needs[next_nodes[v][index - 1]]:
                needs[v] = True

        while index < len(next_nodes[v]):
            w = next_nodes[v][index]
            if parents[v] == w:
                index += 1
                continue
                
            parents[w] = v
            stack.append((v, index + 1))
            stack.append((w, 0))
            break
    new_next_nodes = [[] for _ in range(N)]
    total_movements = 0
    for i in range(N):
        if needs[i]:
            for w in next_nodes[i]:
                if needs[w]:
                    new_next_nodes[i].append(w)
                    total_movements += 1
    
    dists = solve(N, new_next_nodes, D[0])
    for i in range(N):
        if needs[i]:
            dists[i] = total_movements - dists[i]

    d_list = []
    for i in range(N):
        if needs[i]:
            d_list.append((i, dists[i]))
    m_dists = bfs(N, next_nodes, d_list)

    for i in range(N):
        if not needs[i]:
            dists[i] = m_dists[i]
    for i in range(N):
        print(dists[i])

    

    
    





if __name__ == "__main__":
    main()
0