結果

問題 No.2337 Equidistant
ユーザー lam6er
提出日時 2025-04-15 21:47:42
言語 PyPy3
(7.3.15)
結果
WA  
実行時間 -
コード長 3,647 bytes
コンパイル時間 204 ms
コンパイル使用メモリ 82,304 KB
実行使用メモリ 421,132 KB
最終ジャッジ日時 2025-04-15 21:49:13
合計ジャッジ時間 26,833 ms
ジャッジサーバーID
(参考情報)
judge5 / judge4
このコードへのチャレンジ
(要ログイン)
ファイルパターン 結果
sample AC * 1
other AC * 5 WA * 16 TLE * 1 -- * 6
権限があれば一括ダウンロードができます

ソースコード

diff #

import sys
from sys import stdin
from collections import defaultdict, deque

sys.setrecursionlimit(1 << 25)

def main():
    input = sys.stdin.read().split()
    ptr = 0
    N, Q = int(input[ptr]), int(input[ptr+1])
    ptr +=2

    adj = [[] for _ in range(N+1)]
    for _ in range(N-1):
        a = int(input[ptr])
        b = int(input[ptr+1])
        adj[a].append(b)
        adj[b].append(a)
        ptr +=2

    LOG = 20
    parent = [[-1]*(N+1) for _ in range(LOG)]
    depth = [0]*(N+1)

    queue = deque([1])
    visited = [False]*(N+1)
    visited[1] = True
    parent[0][1] = -1

    while queue:
        u = queue.popleft()
        for v in adj[u]:
            if not visited[v] and v != parent[0][u]:
                parent[0][v] = u
                depth[v] = depth[u] + 1
                visited[v] = True
                queue.append(v)

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

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

    def get_kth_ancestor(u, k):
        if k == 0:
            return u
        for i in range(LOG):
            if k & (1 << i):
                u = parent[i][u]
                if u == -1:
                    return -1
        return u

    size = [1]*(N+1)
    visited = [False]*(N+1)
    def post_order(u):
        visited[u] = True
        for v in adj[u]:
            if not visited[v] and v != parent[0][u]:
                post_order(v)
                size[u] += size[v]
    post_order(1)

    size_map = [defaultdict(int) for _ in range(N+1)]
    for u in range(1, N+1):
        for v in adj[u]:
            if parent[0][v] == u:
                size_map[u][v] = size[v]
            else:
                size_map[u][v] = N - size[u]

    for _ in range(Q):
        S = int(input[ptr])
        T = int(input[ptr+1])
        ptr +=2

        if S == T:
            print(N)
            continue

        L = lca(S, T)
        a = depth[S] - depth[L]
        b = depth[T] - depth[L]
        D = a + b

        if D % 2 != 0:
            print(0)
            continue

        k = D // 2

        if k <= a:
            M = get_kth_ancestor(S, k)
            steps_B = (b - (k - a)) -1
            if steps_B < 0:
                B = T
            else:
                B = get_kth_ancestor(T, steps_B)
            A = get_kth_ancestor(S, k-1)
        else:
            steps_from_T = b - (k - a)
            M = get_kth_ancestor(T, steps_from_T)
            A = get_kth_ancestor(S, k-1)
            steps_B = steps_from_T -1
            if steps_B < 0:
                B = T
            else:
                B = get_kth_ancestor(T, steps_B)

        found_A = False
        found_B = False
        size_A = 0
        size_B = 0
        for v in adj[M]:
            if v == A:
                size_A = size_map[M][v]
                found_A = True
            if v == B:
                size_B = size_map[M][v]
                found_B = True
        if not found_A or not found_B:
            print(0)
            continue

        ans = N - size_A - size_B
        print(ans)

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