백준 11438 - LCA 2 (Python)

전처리로 부모 관계를 정해준 뒤 푸는게 핵심입니다.
pypy3로 제출해야 합니다.

전체 코드

import sys

sys.setrecursionlimit(100000)
input = sys.stdin.readline
LOG = 21

n = int(input())
parent = [[0] * LOG for _ in range(n + 1)]
visited = [0] * (n + 1)
distance = [0] * (n + 1)
graph = [[] for _ in range(n + 1)]

for _ in range(n - 1):
    a, b = map(int, input().split())
    graph[a].append(b)
    graph[b].append(a)

def dfs(current, depth):
    visited[current] = True
    distance[current] = depth

    for next in graph[current]:
        if visited[next]:
            continue
            
        parent[next][0] = current
        dfs(next, depth + 1)

def set_parent():
    dfs(1, 0)
    for i in range(1, LOG):
        for j in range(1, n + 1):
            parent[j][i] = parent[parent[j][i - 1]][i - 1]


def lca(a, b):
    if distance[a] < distance[b]:
        a, b = b, a

    for i in range(LOG - 1, -1, -1):
        if distance[a] - distance[b] >= 2**i:
            a = parent[a][i]

    if a == b:
        return a

    for i in range(LOG - 1, -1, -1):
        if parent[a][i] != parent[b][i]:
            a = parent[a][i]
            b = parent[b][i]

    return parent[a][0]

set_parent()
m = int(input())

for _ in range(m):
    a, b = map(int, input().split())
    print(lca(a, b))