백준(BOJ) 19581 두 번째 트리의 지름 파이썬(Python)

문제 링크

트리 지름 문제의 변형판 문제.
우선 트리의 지름을 구하는 알고리즘을 떠올려보자.

  1. 우선 아무 정점에서 시작해 DFS로 가장 먼 정점을 찾는다.
  2. 찾은 정점에서 DFS를 한번 더 이용해 가장 먼 정점을 찾는다.
  3. 찾은 두 정점 사이의 거리가 트리에서 가장 먼 거리, 즉 지름이 된다.
    이 문제를 풀기 위해서 트리 지름을 구한 이후에, 지름을 구성하는 두 노드를 각각 제외한 두 개의 트리에서 지름을 구해 그 중 최댓값을 취하는 방법을 썼다.
    이를테면 어떤 트리 \(T\)에서, 트리의 지름을 구성하는 두 노드를 \(v\)와 \(w\)라고 한다면, 그 두 노드는 당연히 리프 노드일 것이다.
    따라서 한번은 \(v\)를 빼고, 다른 한번은 \(w\)를 빼고 다시 지름을 구한다면 그 중 큰 값이 두 번째 지름일 것이다.
    왜냐하면, 트리의 지름 문제는 결국 노드(구체적으로는 리프 노드) 두 개의 쌍을 선택해서, 거리가 가장 큰 것을 고르는 문제고, 두 번째 지름을 구하는 문제는 모든 노드 두 개의 쌍 중 트리의 지름이 되는 쌍을 제외한 쌍 중 거리가 가장 큰 것을 구하는 문제이기 때문이다.
    처음에는 DFS를 재귀로 구현했는데, 파이썬의 재귀가 끔찍하게 느리다 보니 TLE를 면치 못했다.
    스택과 반복문을 사용하도록 DFS의 구현을 바꾸자 문제 없이 통과했다.
    참고로, 이것보다 조금 더 빠르게 풀 수도 있는데, 어차피 두 번째 지름도 원래 지름을 이루는 \(v\) 또는 \(w\) 중 하나를 포함해야 한다는 점을 이용하는 것이다.
    \(v\) 또는 \(w\)가 두 번째 지름에 포함된다는 것은 트리의 지름을 구하는 알고리즘을 증명하는 것과 유사한 귀류법으로 쉽게 알 수 있고, 따라서 지름을 세 번 구하는 대신, 지름을 한 번 구한 다음 지름을 이루는 각 노드에서 DFS를 한다면, DFS를 6번이 아니라 4번만 해서 답을 구할 수 있다.
    물론 시간복잡도는 \(O(6N)\)이든 \(O(4N)\)이든 \(O(N)\)이기 때문에 드라마틱한 차이는 아니다.
#!python

from sys import stdin

input = stdin.readline

# dfs를 이용해 출발점인 node에서의 최장거리와 그 거리를 지닌 노드를 반환한다
# omit은 거리를 구할 때 사용하지 않을(제외할) 노드를 의미한다
def dfs(node, omit):
    stack = [node]
    distance = [None] * (n + 1)
    distance[node] = 0
    distance[omit] = -1
    max_dist, dist_node = 0, None
    while stack:
        curr_node = stack.pop()
        for succ_node, succ_weight in tree[curr_node]:
            if distance[succ_node] is not None:
                continue
            distance[succ_node] = distance[curr_node] + succ_weight
            if max_dist < distance[succ_node]:
                max_dist = distance[succ_node]
                dist_node = succ_node
            stack.append(succ_node)
    return max_dist, dist_node

# dfs 두 번으로 지름을 구한다
# omit은 지름을 구할 때 사용하지 않을(제외할) 노드를 의미한다
# 트리의 지름과 지름을 이루는 두 노드를 반환한다
def get_diameter(omit):
    _, start_node = dfs(1, omit)
    diameter, end_node = dfs(start_node, omit)
    return diameter, start_node, end_node


n = int(input())
tree = [[] for _ in range(n + 1)]

for _ in range(n - 1):
    src, dst, weight = map(int, input().split())
    tree[src].append((dst, weight))
    tree[dst].append((src, weight))

# 0번 노드는 어차피 사용하지 않으므로 omit=0은 모든 노드를 사용함을 의미
_, s, e = get_diameter(0)

# 지름을 이루는 두 노드를 사용하지 않고 지름을 구해서, 더 큰 쪽을 취한다
print(max(get_diameter(s)[0], get_diameter(e)[0]))

2024

맨 위로 이동 ↑

2023

세그먼트 트리

개요 선형적인 자료구조에서는 값에 접근하는 데에 \(O(1)\)이면 충분하지만, 대신 부분합을 구하는 데에는 \(O(N)\)이 필요하다. 그렇다면 이 자료구조를 이진 트리로 구성하면 어떨까? 값에 접근하는 데에 걸리는 시간이 \(O(\lg N)\)으로 늘어나지만 대신 부분합을 구하...

벨만-포드 알고리즘

개요 다익스트라 알고리즘과 함께 Single Sourse Shortest Path(SSSP) 문제를 푸는 알고리즘이다. 즉, 한 노드에서 다른 모든 노드로 가는 최단 경로를 구하는 알고리즘이다. 다익스트라 알고리즘보다 느리지만, 음수 가중치 간선이 있어도 작동하며, 음수 가중치 사...

다익스트라 알고리즘

개요 다익스트라 알고리즘은 Single Sourse Shortest Path(SSSP) 문제를 푸는 알고리즘 중 하나이다. 즉, 한 노드에서 다른 모든 노드로 가는 최단 경로를 구하는 알고리즘이다. 단, 다익스트라 알고리즘은 음수 가중치 엣지를 허용하지 않는다. 이 경우에는 벨만-...

맨 위로 이동 ↑