트리 트리오 중간값

• 구하라고 하는 값이 설명으로는 좀 복잡해보이는데, 따져보면 단순하다.
• 우선, 트리의 지름이란 트리에 존재하는 경로중 가장 긴 경로의 길이를 의미한다는 것을 기억하자.
• v1과 v2 사이의 거리가 트리의 지름일 때 이 값을 d라고 하면,
  • v1과 v3 사이의 거리 또는 v2과 v3 사이의 거리가 d인 v3가 존재하면, 답은 d이다.
    • (v1, v2, v3)를 잡으면 되니까..
  • 그런 v3가 존재하지 않는다면 답은 d-1 이다
    • 1) 이 경우, 중간값이 d인 세 점은 존재할 수 없다. 만약 다른 두 노드 v3, v4사이의 거리가 d라면, 즉 dist(v3,v4)=d라면, 역시 dist(v1,v3)=dist(v2,v3)=d 도 성립하게 되기 때문이다.
    • 2) 중간값이 d-1이 되는 세 점은 언제나 존재한다. v3를 v2의 이웃 노드로 잡으면 dist(v1,v2)=d, dist(v1,v3)=d-1, dist(v2,v3)=1 이니까 중간값은 d-1이 된다
• 구현은 위 로직을 그대로 따르면 된다
  • 1) 트리의 지름을 구하기 위해, 임의의 노드로부터 가장 먼 노드 v1을 구한다
  • 2) v1으로부터 거리가 가장 먼 노드 v2를 구한다.
  • 3) 거리가 가장 먼 노드가 2개 이상 존재하면 그 거리(=지름) d를 리턴
  • 4) v2으로부터 거리가 d인 노드가 2개 이상 존재하면 d를 리턴
  • 5) 다 해당 안되면 d-1을 리턴
• 3번의 BFS로 해결 가능하다. 시간 복잡도는 3*O(n) = O(n)

"""Solution code for "Programmers 68937. 트리 트리오 중간값".

- Problem link: https://programmers.co.kr/learn/courses/30/lessons/68937
- Solution link: http://www.teferi.net/ps/problems/programmers/68937
"""

import operator
from teflib import tgraph


def solution(n, edges):
    graph = [set() for _ in range(n)]
    for u, v in edges:
        graph[u - 1].add(v - 1)
        graph[v - 1].add(u - 1)

    dists = tgraph.min_distances(graph, 0)
    v1 = max(enumerate(dists), key=operator.itemgetter(1))[0]
    dists_from_v1 = tgraph.min_distances(graph, v1)
    diameter = max(dists_from_v1)
    v2 = [u for u, dist in enumerate(dists_from_v1) if dist == diameter]
    if len(v2) > 1 or tgraph.min_distances(graph, v2[0]).count(diameter) > 1:
        return diameter
    return diameter - 1