백준 1167 트리의 지름 파이썬

dasd412·2022년 5월 20일
0

코딩 테스트

목록 보기
39/71

문제 설명

트리의 지름이란, 트리에서 임의의 두 점 사이의 거리 중 가장 긴 것을 말한다. 트리의 지름을 구하는 프로그램을 작성하시오.

입력 조건

트리가 입력으로 주어진다. 먼저 첫 번째 줄에서는 트리의 정점의 개수 V가 주어지고 (2 ≤ V ≤ 100,000)둘째 줄부터 V개의 줄에 걸쳐 간선의 정보가 다음과 같이 주어진다. 정점 번호는 1부터 V까지 매겨져 있다.

먼저 정점 번호가 주어지고, 이어서 연결된 간선의 정보를 의미하는 정수가 두 개씩 주어지는데, 하나는 정점번호, 다른 하나는 그 정점까지의 거리이다. 예를 들어 네 번째 줄의 경우 정점 3은 정점 1과 거리가 2인 간선으로 연결되어 있고, 정점 4와는 거리가 3인 간선으로 연결되어 있는 것을 보여준다. 각 줄의 마지막에는 -1이 입력으로 주어진다. 주어지는 거리는 모두 10,000 이하의 자연수이다.

출력 조건

첫째 줄에 트리의 지름을 출력한다.


전체 코드

import sys

sys.setrecursionlimit(10 ** 5)

v = int(sys.stdin.readline().rstrip())

graph = dict()
for i in range(v):
    data = list(map(int, sys.stdin.readline().split()))
    # graph[출발지]=[(목적지,비용)]
    graph[data[0]] = []
    for j in range(1, len(data) - 1, 2):
        graph[data[0]].append((data[j], data[j + 1]))

# 양방향처럼 되어 있기 때문에 방문 여부 상태는 저장할 필요가 있다. 안그러면 무한 재귀 발생
visited = set()

answer = 0


# current : 현재 노드 번호
def get_length_between_nodes(current: int) -> int:
    global answer

    # 현재 노드에서 연결된 '길이' 들을 뜻한다.
    lengths = []

    # 먼저 현재 노드와 이어진 곳들에 대해 알아본다.
    for adj_node, adj_cost in graph[current]:
        # 인접한 노드가 방문한 적 없다면,
        if adj_node not in visited:
            # 재 방문을 막기 위해 방문 표시한다.
            visited.add(adj_node)

            # 메시지 전달로 얻은 리턴 값들은 하나의 노드에서 이어진 길이들을 뜻한다.
            length = adj_cost + get_length_between_nodes(adj_node)
            lengths.append(length)

    # 만약 리프노드라면, 전부 방문해봤으므로 lengths 의 길이가 0이다.
    if len(lengths) == 0:
        return 0
    # 이어진 게 하나 있으면 해당 길이 리턴
    elif len(lengths) == 1:
        answer = max(answer, lengths[0])
        return lengths[0]
    # 이어진 것이 여러가지라면, 길이들을 내림 차순 정렬한다.
    else:
        lengths.sort(reverse=True)
        # 이어진 것 중 내림차순으로 긴 것 2개를 뽑아서 더해본다. 최대 길이보다 길다면 갱신
        answer = max(answer, lengths[0] + lengths[1])
        # 가장 긴 길이를 리턴한다.
        return lengths[0]


# 루트는 그냥 1이라고 가정한다. 2<=v이기 때문.
visited.add(1)
get_length_between_nodes(1)

print(answer)

해설

예시 입력은 다음과 같다. 답은 20이 되야 한다.

6
1 2 3 -1
2 1 3 5 3 3 5 -1
3 2 5 4 7 -1
4 3 7 -1
5 2 3 6 5 -1
6 5 5 -1

위 코드에서 lengths 배열의 길이가 2이상인 부분만 설명한다.
원하는 답은 트리의 지름이다. 따라서 어떤 노드와 연결된 간선이 무수히 많더라도, 지름의 특성 상 어떤 노드와 연결된 간선 중 2개만 선택하면 된다. 즉, 연결된 간선들 중 길이가 긴 순으로 2개만 선택하면, 해당 노드에서의 지름을 계산할 수 있다.


참고

https://my-coding-notes.tistory.com/286
https://tsgoing.tistory.com/28

profile
아키텍쳐 설계와 테스트 코드에 관심이 많음.

0개의 댓글