[트리 트리오 중간값/LV4/Python3]

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

 

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번 정점 사이에 간선이 있음을 의미합니다.
1 <= v1, v2 <= n, v1!= v2
입력으로 주어지는 그래프는 항상 트리입니다.


테스트 케이스는 문제 풀이 당시 사용했던 케이스도 추가되어 있습니다.

n	edges	result
4	[[1,2], [2,3], [3,4]]	2
5	[[1,5], [2,5], [3,5], [4,5]]	2
8	[[1,2], [2,3], [3,4], [4,5], [3,6], [6,7], [7,8]]	5
7	[[1,2], [2,3], [3,4], [4,5], [3,6], [6,7]]	4

 

 

백준에서 풀었던 트리의 지름 문제에서 아이디어를 얻었다.

 

지름을 구하는 방식은 다음과 같다.

특정 정점 v에서 가장 거리가 먼 정점을 u를 구한다.

정점 u에서 가장 거리가 먼 정점 w를 구한다.

이때 u와 w 사이의 거리가 트리에서의 최장 거리이며, 지름이 된다.

 

 

위 문제의 경우 세 정점 사이의 거리 중 중앙값을 구하는 문제이다.

이때 중앙값의 경우 항상 (지름의 길이 -1) 혹은 (지름의 길이)를 보장하게 된다는 것을 가정하고 문제를 풀었다.

 

 

정점 a와 b가 트리의 지름일 때 c를 선택할 수 있는 경우는 2가지가 존재한다.

1. a와 b사이의 정점 c를 선택한 경우

2. a와 b사이가 아닌 정점 c를 선택한 경우

 

1의 경우에는 a나 b 바로 앞의 정점을 선택할 수 있다.

a와 b 사이의 거리가 d라면 각각의 거리를 구하면 1, d-1, d로 구할 수 있다. 이때의 경우 중앙값은 항상 (지름 - 1)이 된다.

 

 

2의 경우는 a와 b로부터 지름만큼 떨어진 곳에 c가 있는지 확인해야 한다.

a와 b 사이의 거리가 d라고 했을 때 a와 c 혹은 b와 c의 거리가 d일 수 있기 때문이다.

 

해당 문제는 bfs로 중앙값을 구할 때 처리하는 방식을 사용했다.

정점 a 혹은 b에서 탐색을 시작한다면, 마지막에 확인되는 정점에는 항상 b나 a가 포함되기 때문이다.

조건식으로 수정한다면 다음과 같게 얻을 수 있다.

(동일한 거리에 정점이 2개 이상 존재하고, 더 이상 탐색이 불가능할 때) == (지름에 해당하는 거리에 정점이 2개 이상 존재한다)를 뜻하게 되므로, 이때는 중앙값을 지름의 길이를 반환하도록 하였다.

 

내가 사용했던 bfs로 중앙값을 구할 때는 특정 정점 1개에 대해서 구하는 것이기 때문에, 2번 경우에는 어떤 정점에서 시작했는지에 따라 답이 나뉘게 되는 예외 케이스가 있다.

 

아래 테스트 케이스 그림의 경우, 정점 1에서 검색을 하면 중앙값 4를 반환하지만, 정점 8에서 시작하면 5를 반환한다.

해당 문제를 해결하기 위해, 트리의 지름이 되는 두 정점 모두에게서 중앙값을 구하고, 그중 더 큰 값을 반환하는 것으로 하였다.

 

 

 

 

from collections import deque

def solution(n, edges):
    #답은 항상 지름 or 지름-1
    #특정 정점을 기준으로 가장 먼 정점 구하기
    
    routes = [[] for _ in range(n+1)]
    for u, v in edges:
        routes[u].append(v)
        routes[v].append(u)

    def find_line(i):
        visit = [False] * (n+1)
        pos = deque()
        pos.append(i)
        visit[i] = True
        while pos:
            v = pos.popleft()
            for vv in routes[v]:
                if visit[vv]: continue
                pos.append(vv)
                visit[vv] = True
        return v


    def find_longest(v):
        visit = [False] * (n+1)
        pos = [v]
        visit[v] = True
        cnt = 0
        while pos:
            next_pos = []
            for i in pos:
                for vv in routes[i]:
                    if visit[vv]: continue
                    visit[vv] = True
                    next_pos.append(vv)

            if not next_pos:
                if len(pos) >= 2:
                    return cnt
                else:
                    return cnt-1

            cnt += 1
            pos = next_pos

    v = find_line(1)
    return max(find_longest(v), find_longest(find_line(v)))

 

 

 

---

 

만일 이 문제에 가중치를 포함한다면 visit확인이 아닌, distance로 변경을 해서 구하면 동일하게 결과는 얻을 수 있을 것 같다.

후에 유사한 문제를 확인해 보고 풀어봐야겠다.