Programmers. 트리 트리오 중간값

읽기 전

• 불필요한 코드나 잘못 작성된 내용에 대한 지적은 언제나 환영합니다.
• 개인적으로 배운 점을 정리한 글입니다.

문제 링크

Programmers. 트리 트리오 중간값

문제 풀이

abc 노드가 있으면 a-b. b-c. c-a길이 중 중간 값을 반환해야 하는데 두 노드가 트리에서 가장 긴 겨오라면 다른 한 노드는 또 다른 가장 긴 경로에 있거나 아니면 가장 긴 경로 사이에 있을 것이다. 따라서 임의의 노드에 대해 가장 긴 경로를 갖는 노드를 찾고 해당 노드에서 가장 긴 경로를 갖는 노드들을 탐색한다. 만약 여러 개라면 중간 값은 가장 긴 경로일 것이다. 그런데 하나라면 그 노드에서 다시 가장 긴 경로를 갖는 노드를 탐색해야 한다. 왜냐하면 처음 최장 경로 탐색 지점은 해당 경로의 한 노드에 불과하므로 다른 방향의 끝점에서도 탐색해야 한다. 그 탐색 결과 여전히 하나라면 최장경로 - 1을 반환한다.

• 입력받은 edges 정보를 hash구조에 넣어서 그래프 완성
• 임의의 점을 기준으로 다익스트라를 끝까지 수행해 임의의 점에서 가장 멀리 있는 점 구함
• 앞서 구한 점이 리프노드임을 보장하므로 다시 다익스트라를 끝까지 수행해 현재 리프노드에서 가장 멀리 떨어진 리프노드들을 구한다. 여러 개면 최장 경로가 2개 이상임을 의미하므로 그대로 최장 경로 길이 리턴, 1개라면 구한 노드에서 반대쪽 리프 노드가 여러 개인지 탐색을 위해 다시 다익스트라를 끝까지 수행
• 수행 결과가 여러 개라면 처음의 리프노드 주위로 다른 리프노드가 있었음을 의미하며 그대로 최장 경로 길이 리턴, 여전히 1개라면 반드시 나머지 하나의 노드는 최장 경로 위에 있음을 의미하므로 최장 경로 길이 - 1을 반환

python 코드

from collections import defaultdict
from heapq import heappop, heappush
def solution(n, edges):
    board = defaultdict(list)
    for srt, dst in edges:
        board[srt].append(dst)
        board[dst].append(srt)
    def compute(node):
        visited = [-1 for _ in range(n + 1)]
        M = 0
        q = [(0, node)]
        visited[node] = 0
        res = []
        while q:
            cost, dst = heappop(q)
            if M == cost:
                res.append(dst)
            elif M < cost:
                res = [dst]
                M = cost
            for v in board[dst]:
                if visited[v] == -1:
                    visited[v] = cost + 1
                    heappush(q, (cost + 1, v))
        return res, M

    visited = [-1 for _ in range(n + 1)]
    visited[edges[0][0]] = 0
    q = [(0, edges[0][0])]
    res = edges[0][0]
    M = 0
    while q:
        cost, dst = heappop(q)
        if M < cost:
            M = cost
            res = dst
        for v in board[dst]:
            if visited[v] == -1:
                visited[v] = cost + 1
                heappush(q, (cost + 1, v))
    res, cost = compute(res)
    if len(res) > 1:
        return cost
    else:
        res, cost = compute(res[0])
        if len(res) == 1:
            return cost - 1
        else:
            return cost