(Python/플5)백준알고리즘 11438번 :LCA2

문제 바로가기

 

문제:

N(2 ≤ N ≤ 100,000)개의 정점으로 이루어진 트리가 주어진다. 트리의 각 정점은 1번부터 N번까지 번호가 매겨져 있으며, 루트는 1번이다. 두 노드의 쌍 M(1 ≤ M ≤ 100,000)개가 주어졌을 때, 두 노드의 가장 가까운 공통 조상이 몇 번인지 출력한다.

입력:

첫째 줄에 노드의 개수 N이 주어지고, 다음 N-1개 줄에는 트리 상에서 연결된 두 정점이 주어진다. 그 다음 줄에는 가장 가까운 공통 조상을 알고싶은 쌍의 개수 M이 주어지고, 다음 M개 줄에는 정점 쌍이 주어진다.

출력:

M개의 줄에 차례대로 입력받은 두 정점의 가장 가까운 공통 조상을 출력한다.

 

일단 그림으로 나타내면 다음과 같다.

 

 

첫 시도

import sys

sys.stdin = open('/Users/song/Desktop/Python/Python/h.txt', 'r')
sys.setrecursionlimit(10**6)
n = int(input())

graph = [[] for _ in range(n + 1)]
# 노드
for i in range(n - 1):
    a, b = map(int, sys.stdin.readline().split(" "))
    graph[a].append(b)
    graph[b].append(a)

parent = [0] * (n + 1)  # 각 노드의 부모 노드 정보
depth = [0] * (n + 1)  # 각 노드까지의 깊이
visited = [0] * (n + 1)  # 방문 여부


def dfs(node, idx):
    visited[node] = 1  # 방문처리해주고
    depth[node] = idx  # 방문했으니 현재위치의 depth를 idx로 해줌 기본은 0
    for next_node in graph[node]:  # 그 다음 방문위치
        if visited[next_node] == 1:  # 이미 방문했으면
            continue  # 넘어가고
        parent[next_node] = node  # 내가 자식들의 부모노드임
        dfs(next_node, idx + 1)  # 깊이 1 추가 가 다음 노드

dfs(1, 0)

m = int(input())


def find_parent(x, y):

    while depth[x] != depth[y]:
        if depth[x] > depth[y]:
            x = parent[x]
        else:
            y = parent[y]
    while x != y:
        x = parent[x]
        y = parent[y]

    return x


for i in range(m):
    x, y = map(int, sys.stdin.readline().split(" "))
    print(find_parent(x, y))

 

메모리 초과 난다.

 

깊이를 맞출때와, 공통조상을 찾을때 한 칸씩 올라가는게 아닌 2^n 으로 올라가는게 포인트다.

import sys

sys.stdin = open('/Users/song/Desktop/Python/Python/h.txt', 'r')
sys.setrecursionlimit(10**6)
input = sys.stdin.readline
LENGTH = 21

n = int(input())
parent = [[0] * LENGTH for _ in range(n + 1)]
visited = [False] * (n + 1)
d = [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(x, depth):
    visited[x] = True
    d[x] = depth

    for node in graph[x]:
        if visited[node]:
            continue
        # 바로 한 단계 위에 부모
        parent[node][0] = x
        dfs(node, depth + 1)

def set_parent():
    # 한 단계 부모 갱신
    dfs(1, 0)
    for i in range(1, LENGTH):
        for j in range(1, n + 1):
            # 각 노드에 대해 2**i번째 부모 정보 갱신
            parent[j][i] = parent[parent[j][i - 1]][i - 1]

def lca(a, b):
    # 무조건 b의 깊이가 더 깊도록 설정
    if d[a] > d[b]:
        a, b = b, a
    # a와 b의 깊이가 동일해주도록 설정
    for i in range(LENGTH - 1, -1, -1):
        if d[b] - d[a] >= 2**i:
            b = parent[b][i]

    if a == b:
        return a

    # 올라가면서 공통 조상 찾기
    for i in range(LENGTH - 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))

 

 

# j 노드의 2**i 번째 노드를 저장한다.
# parent[j][i] = parent[parent[j][i - 1]][i - 1]