파이썬 코딩테스트 강좌, 트리의 지름 구하기

트리는 컴퓨터 과학에서 매우 중요한 자료구조 중 하나입니다. 특히 트리는 계층적 관계를 나타내는 데 유용하며, 다양한 알고리즘 문제에서 사용됩니다. 이번 강좌에서는 트리의 지름을 구하는 문제를 다룰 것입니다.
지름이란 트리에서 두 노드 사이의 가장 긴 경로를 의미합니다. 이 문제는 DFS(깊이 우선 탐색) 또는 BFS(너비 우선 탐색) 알고리즘을 통해 해결할 수 있습니다.

문제 설명

주어진 비어있지 않은 트리의 각 노드는 정수로 표시됩니다. 트리에서 두 노드 사이의 가장 긴 경로의 길이를 구하는 문제를 해결하세요.
입력으로는 트리의 노드 수 N과 N-1개의 간선 정보가 주어집니다. 간선 정보는 두 노드를 연결하는 형태로 제공됩니다.
구체적으로, 다음과 같은 형식의 입력이 주어질 것입니다:

    N
    a1 b1
    a2 b2
    ...
    a(N-1) b(N-1)

여기서 a와 b는 각각 연결된 두 노드를 나타냅니다.

입력 예시

출력 예시

이 경우 트리의 지름은 노드 4와 노드 5 사이로, 경로는 4 → 2 → 1 → 3 또는 4 → 2 → 5로 구성됩니다.
따라서 출력은 3입니다.

문제 풀이

트리의 지름을 구하기 위해서는 DFS 또는 BFS 알고리즘을 사용할 수 있습니다.
일반적인 접근 방법은 다음과 같습니다.

첫 번째 단계로, 임의의 노드에서 DFS를 실행해서 가장 먼 노드를 찾습니다.
이 노드를 X라고 합시다.
두 번째 단계로, 노드 X에서 다시 DFS를 실행하여 가장 먼 노드인 Y를 찾습니다.
X와 Y 사이의 경로가 트리의 지름이 됩니다.

이 과정을 통해 시간 복잡도는 O(N)이 되며, 재귀적으로 DFS를 호출하는 방식으로 구현됩니다.

파이썬 코드 구현

이제 위의 로직을 바탕으로 파이썬에서 트리의 지름을 구하는 코드를 구현해 보겠습니다.
아랫부분에 제공하는 코드와 함께 각 단계의 세부 과정을 점검해 보세요.

from collections import defaultdict

def dfs(graph, node, visited):
    visited.add(node)
    max_distance = 0
    farthest_node = node

    for neighbor in graph[node]:
        if neighbor not in visited:
            distance, next_node = dfs(graph, neighbor, visited)
            distance += 1
            
            if distance > max_distance:
                max_distance = distance
                farthest_node = next_node

    return max_distance, farthest_node

def tree_diameter(edges, n):
    graph = defaultdict(list)
    
    for a, b in edges:
        graph[a].append(b)
        graph[b].append(a)

    # Step 1: start DFS from an arbitrary node (1)
    visited = set()
    _, farthest_node = dfs(graph, 1, visited)

    # Step 2: start DFS from the farthest node found
    visited.clear()
    diameter, _ = dfs(graph, farthest_node, visited)

    return diameter

# Input reading part
n = int(input())
edges = [tuple(map(int, input().split())) for _ in range(n-1)]
print(tree_diameter(edges, n))

코드 설명

위의 코드는 다음과 같은 방식으로 구성되어 있습니다:

collections.defaultdict를 사용하여 그래프를 인접 리스트 형태로 만듭니다.
이 때 노드끼리의 연결관계를 나타냅니다.
dfs 함수는 깊이 우선 탐색을 수행하며, 각 노드까지의 거리를 계산합니다.
가장 먼 노드와 거리를 반환합니다.
tree_diameter 함수는 전체적인 과정을 조정하고 두 번의 DFS 호출로 지름을 계산합니다.
마지막 부분에서 사용자로부터 입력을 받고, tree_diameter 함수를 호출하여 결과를 출력합니다.

성능 분석

제시된 알고리즘은 O(N)의 시간 복잡도를 갖고 있습니다.
이는 트리의 모든 노드를 DFS로 한 번씩 방문하기 때문에 가능합니다.
따라서 매우 큰 트리에서도 효율적으로 지름을 계산할 수 있습니다.

결론

이번 강좌에서는 트리의 지름에 대해 살펴보았습니다.
DFS를 이용한 접근법으로 효율적으로 문제를 해결할 수 있었습니다.
트리는 다양한 문제에서 활용되기 때문에 이번 강좌의 내용을 잘 익혀두면 좋습니다.
추가적인 질문이나 연습문제가 필요하다면 댓글로 남겨주세요.