[프로그래머스] 트리 트리오 중간값

Posted on

문제 링크 : https://programmers.co.kr/learn/courses/30/lessons/68937#_=_

프로그래머스 - 트리 트리오 중간값

문제 설명

n개의 점으로 이루어진 트리가 있습니다. 이때, 트리 상에서 다음과 같은 것들을 정의합니다.

  • 어떤 두 점 사이의 거리는, 두 점을 잇는 경로 상 간선의 개수로 정의합니다.
  • 임의의 3개의 점 a, b, c에 대한 함수 f(a, b, c)의 값을 a와 b 사이의 거리, b와 c 사이의 거리, c와 a 사이의 거리, 3개 값의 중간값으로 정의합니다.

트리의 정점의 개수 n과 트리의 간선을 나타내는 2차원 정수 배열 edges가 매개변수로 주어집니다. 주어진 트리에서 임의의 3개의 점을 뽑아 만들 수 있는 모든 f값 중에서, 제일 큰 값을 구해 return 하도록 solution 함수를 완성해주세요.

제한 사항

  • n은 3 이상 250,000 이하입니다.
  • edges의 행의 개수는 n-1 입니다.
    • edges의 각 행은 [v1, v2] 2개의 정수로 이루어져 있으며, 이는 v1번 정점과 v2번 정점 사이에 간선이 있음을 의미합니다.
    • v1, v2는 각각 1 이상 n 이하입니다.
    • v1, v2는 다른 수입니다.
    • 입력으로 주어지는 그래프는 항상 트리입니다.

처음에는 이 문제를 정말 단순하게 생각했다. 모든 점들의 거리를 구한 후 직접 조합을 이용해서 사용하면 되지 않을까 했었는데.. 그렇게 하면 타임아웃 에러가 난다.

그렇기 때문에 간단한 아이디어를 하나 접목시켜주어야 했는데, 그것은 바로 아래와 같다.

  1. BFS를 통해 특정 s 위치에서 각 지점 사이의 거리를 계산한다
  2. 최댓값을 d라고 할 때, 최댓값의 갯수가 2개 이상일때는 d를 리턴, 1개일 때에는 d-1을 리턴한다.

이것이 가능한 이유는 바로 “중간값”이라는 개념 때문이다. 최대값이 d일 때 가장 큰 중간값이 될 수 있는 수는 d 혹은 d-1이다. 이 개념을 이용하면 되었던 문제!

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
from collections import deque

def BFS(tree, start):
    distance = []
    q = deque()
    q.append([start, 0]) # 시작 위치와 거리!
    visit = {start : True} # 시작 위치를 방문했었으면 True라고 입력해준다.

    while q:
        curr, dist = q.popleft()
        distance.append([curr, dist])

        for node in tree[curr]:
            if node not in visit:
                q.append([node, dist+1])
                visit[node] = True

    return distance

def solution(n, edges):
    tree = {i+1 : [] for i in range(n)}
    for edge in edges:
        tree[edge[0]].append(edge[1])
        tree[edge[1]].append(edge[0])

    # 루트에서 가장 멀리 떨어진 노드 탐색    
    leaf = BFS(tree, 1)
    # 루트에서 가장 멀리 떨어진 노드에서 거리 탐색
    farther = BFS(tree, leaf[-1][0])

    # 거리가 가장 먼 leaf가 두개 이상일 경우
    if leaf[-1][1] == leaf[-2][1]:
        return farther[-1][1]
    else:
        # 해당 리프에서 두번째로 멀리 떨어진 길이가 중간값. 
        return farther[-2][1]

Reference