백준 문제 풀이 - 트리의 지름 1967번

📜 문제 이해하기

시간 제한 메모리 제한 제출 정답 맞힌 사람 정답 비율
2 초 128 MB 24227 9784 7416 43.124%
문제
트리(tree)는 사이클이 없는 무방향 그래프이다. 트리에서는 어떤 두 노드를 선택해도 둘 사이에 경로가 항상 하나만 존재하게 된다. 트리에서 어떤 두 노드를 선택해서 양쪽으로 쫙 당길 때, 가장 길게 늘어나는 경우가 있을 것이다. 이럴 때 트리의 모든 노드들은 이 두 노드를 지름의 끝 점으로 하는 원 안에 들어가게 된다.

이런 두 노드 사이의 경로의 길이를 트리의 지름이라고 한다. 정확히 정의하자면 트리에 존재하는 모든 경로들 중에서 가장 긴 것의 길이를 말한다.
입력으로 루트가 있는 트리를 가중치가 있는 간선들로 줄 때, 트리의 지름을 구해서 출력하는 프로그램을 작성하시오. 아래와 같은 트리가 주어진다면 트리의 지름은 45가 된다.

트리의 노드는 1부터 n까지 번호가 매겨져 있다.

💡 문제 재정의

각 노드에서 가장 길이가 긴 두개의 경로를 찾고 두 개의 합을 찾자.

✏️ 계획 수립

먼저 무방향 가중치 그래프를 받기 위해 인접 리스트를 구현한다.
그 후 bfs를 진행하는데 1번 노드는 항상 root이기에 1번 노드를 시작으로 bfs를 시작한다.
bfs에는 자신의 자식 노드들 중에서 가장 큰 weight를 반납한다. 이 때 마지막 leaf 노드일때는 0을 반납한다. bfs를 수행하면서 자식 노드들에 대한 weight들을 전부 저장한다.
bfs를 수행하고 나서 자식 노드들에 대한 weight들에 대한 값을 얻으면 가장 긴 2개의 길이의 합이 가장 큰 것을 구하면 된다. 이 때 자식 노드가 1개밖에 없는 경우도 있기 때문에 이때는 그 값이 지름이 된다.

💻 계획 수행

import sys

sys.setrecursionlimit(100000)
adj_list = []
visited = []
node_diameter = []


def bfs(node):
    visited[node] = 1
    # leaf 일때
    if not adj_list[node]:
        return 0

    max_weight = 0
    for next_node in adj_list[node]:
        if not visited[next_node[0]]:
            node_weight = next_node[1] + bfs(next_node[0])
            node_diameter[node].append(node_weight)
            if max_weight < node_weight:
                max_weight = node_weight
    return max_weight


if __name__ == '__main__':
    N = int(sys.stdin.readline().strip())
    if N == 1:
        print(0)
    else:
        visited = [0 for _ in range(N + 1)]
        node_diameter = [[] for _ in range(N + 1)]
        adj_list = [[] for _ in range(N+1)]

        for _ in range(N - 1):
            node1, node2, weight = map(int, sys.stdin.readline().split())
            adj_list[node1].append([node2, weight])
            adj_list[node2].append([node1, weight])

        bfs(1)

        max_diameter = 0
        for weights in node_diameter:
            if len(weights) >= 2:
                first_max = weights.pop(weights.index(max(weights)))
                second_max = weights.pop(weights.index(max(weights)))
                diameter = first_max + second_max
            elif len(weights) == 1:
                diameter = weights.pop()
            else:
                diameter = 0
            if diameter > max_diameter:
                max_diameter = diameter
        print(max_diameter)

🤔 회고

처음에는 leaf노드에서 bfs를 진행하여 가장 긴 경로를 찾는 알고리즘으로 구현했으나 시간 초과가 발생하였고 고민하다가 한 번의 bfs로 수행할 수 있다는 것을 깨닫고 새롭게 구현해보았다. 여러 번의 bfs는 시간 초과가 발생할 수 있는 문제였다.