새싹프로그래머의 이야기

https://www.acmicpc.net/problem/1167

임의의 점으로 부터 가장 먼 노드를 2번 구해 지름을 구한다는 발상에 너무 감탄했던 문제입니다.

문제: 트리가 주어졌을 때, 트리의 지름을 구하는 문제입니다.

풀이 방법: 다잌스트라

풀이 코드:

import random
import sys
import heapq

n = int(input())

graph = [[] for _ in range(n+1)]
dist = [1E9] * (n+1)

for i in range(1,n+1):
    line = list(map(int,sys.stdin.readline().rstrip().split()))
    nowNode = line[0]

    for j in range(1,len(line)-1,2):
        graph[nowNode].append((line[j+1], line[j])) #cost, nextNode

def dijkstra(node):
    q = [(0,node)]

    while q:
        startCost, startNode = heapq.heappop(q)

        for nextCost,nextNode in graph[startNode]:
            distance = dist[startNode] + nextCost

            if dist[nextNode] > distance:
                dist[nextNode] = distance
                heapq.heappush(q,(distance,nextNode))

#임의의 점으로 부터 제일 먼 노드 X 구하기
startpos = random.randint(1,n)
dist[startpos] = 0
dijkstra(startpos)
index = dist.index(max(dist[1:]))

#X로부터 제일 먼 노드 Y 구하기
dist = [1E9] * (n+1)
dist[index] = 0
dijkstra(index)

#X-Y의 거리가 지름이 된다
print(max(dist[1:]))

풀이 설명:

먼저 입력 값을 보면,

1 3 2 -1

과 같이 입력이 주어질 경우, 1에서 3까지 가는 비용이 2가 된다는 의미이고

3 1 2 4 3 -1

과 같이 입력이 주어지면, 3에서 1까지 가는 비용이 2 그리고 3에서 4까지 가는 비용이 3이라는 의미입니다.

for i in range(1,n+1):
    line = list(map(int,sys.stdin.readline().rstrip().split()))
    nowNode = line[0]

    for j in range(1,len(line)-1,2):
        graph[nowNode].append((line[j+1], line[j])) #cost, nextNode

방향성이 없고 한줄에 한 노드로 부터 갈 수 있는 모든 노드의 정보가 주어지므로 다음과 같이 입력을 받아주었습니다.

def dijkstra(node):
    q = [(0,node)]

    while q:
        startCost, startNode = heapq.heappop(q)

        for nextCost,nextNode in graph[startNode]:
            distance = dist[startNode] + nextCost

            if dist[nextNode] > distance:
                dist[nextNode] = distance
                heapq.heappush(q,(distance,nextNode))

다음은 가장 거리가 먼 노드를 구하기 위해 다잌스트라 함수를 구현했습니다.

#임의의 점으로 부터 제일 먼 노드 X 구하기
startpos = random.randint(1,n)
dist[startpos] = 0
dijkstra(startpos)
index = dist.index(max(dist[1:]))

#X로부터 제일 먼 노드 Y 구하기
dist = [1E9] * (n+1)
dist[index] = 0
dijkstra(index)

#X-Y의 거리가 지름이 된다
print(max(dist[1:]))

그런 다음, 임의의점 하나로 부터 제일 먼 노드 X를 구해줍니다.

또, X로 부터 제일 먼 노드 Y를 구해줍니다.

이때, X로부터 Y까지의 거리가 트리의 지름이 되게 됩니다.

이 아이디어 발상에 감탄한 문제였습니다.