문제 요약
트리에서 서로 다른 세 정점 a, b, c를 골랐을 때, 세 쌍의 거리의 중간값을 f(a, b, c)라고 한다.
모든 세 정점 조합 중 f의 최댓값을 구한다.
핵심 아이디어
트리에서 가장 먼 두 정점 사이의 거리를 지름이라고 하자. 어떤 세 정점을 골라도 두 정점 사이의 거리는 지름을 넘을 수 없으므로, f의 최댓값도 지름을 넘을 수 없다.
지름의 양 끝을 u, v, 지름 길이를 D라고 하자.
답이 D가 되는 경우
u에서 거리가 D인 정점이 둘 이상 있다면, 그중 두 정점과 u를 고를 수 있다.
u와 x의 거리 = D
u와 y의 거리 = D
x와 y의 거리 <= D세 거리의 중간값은 D이므로 답은 D다.
u에서는 가장 먼 정점이 하나뿐이어도, v에서 거리가 D인 정점이 둘 이상이라면 같은 이유로 답은 D다.
답이 D - 1이 되는 경우
지름 양 끝 모두에서 가장 먼 정점이 하나뿐이면, 지름 길이 D를 두 번 포함하는 세 정점을 만들 수 없다. 이때 최댓값은 D - 1이 된다.
따라서 다음만 확인하면 된다.
- 지름의 한 끝
u를 찾는다. u에서 가장 먼 정점들의 개수를 센다.- 그 개수가 2 이상이면
D를 반환한다. - 그렇지 않으면 반대 끝
v에서도 같은 검사를 한다. - 둘 다 하나뿐이면
D - 1을 반환한다.
지름의 한 끝 찾기
트리의 임의의 정점에서 BFS를 수행해 가장 먼 정점을 찾으면, 그 정점은 지름의 한 끝이 된다.
그 정점에서 다시 BFS를 수행하면 지름 길이와 반대쪽 끝을 찾을 수 있다.
풀이 과정
- 간선 정보로 인접 리스트를 만든다.
- 1번 정점에서 BFS를 수행해 지름의 한 끝
u를 찾는다. u에서 BFS를 수행해 지름 길이D와 최장 거리 정점 개수를 구한다.- 최장 거리 정점이 2개 이상이면
D를 반환한다. - 반대쪽 지름 끝
v에서 BFS를 수행한다. v에서도 최장 거리 정점이 2개 이상이면D, 아니면D - 1을 반환한다.
Python 코드
from collections import deque
def solution(n, edges):
graph = [[] for _ in range(n)]
for left, right in edges:
left -= 1
right -= 1
graph[left].append(right)
graph[right].append(left)
def bfs(start):
distance = [-1] * n
distance[start] = 0
queue = deque([start])
farthest_node = start
max_distance = 0
farthest_count = 1
while queue:
node = queue.popleft()
for next_node in graph[node]:
if distance[next_node] != -1:
continue
distance[next_node] = distance[node] + 1
queue.append(next_node)
if distance[next_node] > max_distance:
max_distance = distance[next_node]
farthest_node = next_node
farthest_count = 1
elif distance[next_node] == max_distance:
farthest_count += 1
return farthest_node, max_distance, farthest_count
# 임의의 정점에서 가장 먼 정점은 지름의 한 끝이다.
diameter_end, _, _ = bfs(0)
# 지름의 한 끝에서 지름 길이와 최장 거리 정점 개수를 구한다.
other_end, diameter, farthest_count = bfs(diameter_end)
if farthest_count >= 2:
return diameter
# 반대쪽 끝에서도 최장 거리 정점이 여러 개인지 확인한다.
_, _, farthest_count = bfs(other_end)
if farthest_count >= 2:
return diameter
return diameter - 1예시
예시 1: 일직선 트리
1 - 2 - 3 - 4지름은 1과 4 사이의 거리 3이다. 지름 끝인 1에서 가장 먼 정점은 4 하나뿐이고, 4에서도 1 하나뿐이다.
따라서 답은 3 - 1 = 2다.
예시 2: 별 모양 트리
1 2
\ /
5
/ \
3 4어떤 잎 정점에서 시작해도 다른 잎 정점 여러 개가 거리 2로 가장 멀다. 지름 길이 2가 두 번 포함되는 세 정점을 만들 수 있으므로 답은 2다.
시간 복잡도
정점 수를 N이라고 하자.
BFS는 간선과 정점을 각각 한 번씩 방문하므로 O(N)이다. BFS를 세 번 수행한다.
- 시간 복잡도:
O(N) - 공간 복잡도:
O(N)
재귀 DFS 대신 BFS를 사용하므로, 정점 수가 25만 개인 긴 일자 트리에서도 재귀 깊이 제한에 걸리지 않는다.
정리
최댓값은 지름 길이 D 또는 D - 1 중 하나다. 지름 끝에서 최장 거리 정점이 여러 개인지 확인하면, 세 정점이 지름 길이를 두 번 만들 수 있는지 판단할 수 있다.
