프로그래머스 트리 트리오 python

https://programmers.co.kr/learn/courses/30/lessons/68937?language=python3 

코딩테스트 연습 - 트리 트리오 중간값

1. 문제 설명

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

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

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

---

2. 제한 사항

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

 

3. 풀이

트리의 지름을 활용하는 문제로 트리내에 트리의 지름이 되는 정점 쌍이 여러개 있을 수 있는 성질을 이용하면 된다.

from collections import deque
def find_fahrest_node(here, n, tree):
	
    """ 
    here로 부터 가장 먼 노드를 찾는 함수로서 가장 먼 노드의 정보 list를 반환한다.
    노드 정보 리스트는 노드 번호, 거리, 횟수로 구성되어 있다. 
    """
    
    fahrest_node_info = [-1, -1, -1] #node, dis, cnt
    
    visited =[False for i in range(n+1)]
    visited[here] = True
    q =deque()
    q.append((here,0))
    
    while q:
        here, depth = q.popleft()
        
        if depth > fahrest_node_info[1]:
            fahrest_node_info[0] = here
            fahrest_node_info[1] = depth
            fahrest_node_info[2] = 1
            
        elif depth == fahrest_node_info[1]:
            fahrest_node_info[2] += 1

        for there in tree[here]:
            if not visited[there]:
                visited[there] = True
                q.append((there, depth+1))
                
    return fahrest_node_info
    
    
def solution(n, edges):
    answer = 0
    INF = float('inf')
    tree = [[] for i in range(n+1)]
    for edge in edges:
        v1, v2 = edge
        tree[v1].append(v2)
        tree[v2].append(v1)
    
    #가장 먼 노드를 2번에 걸쳐서 찾으면 두 노드사이의 거리가 트리의 지름이 된다.
    v1, d1, _ = find_fahrest_node(1, n, tree)
    v2, d2, cnt2 = find_fahrest_node(v1, n, tree)
    
    #만약 트리의 지름을 갖는 노드가 2개 이상이라면 
    #정답은 트리의 지름가 된다. [트리의 지름, 트리의 지름, 어떤값]의 중간값은 무조건 트리의 지름 
    if cnt2 >= 2:
        return d2
    else: 
    	#만약 트리의 지름의 길이가 1개라면 가장 먼 노드를 한번 더 해야하는데
        #그 이유는 시작노드가 루트노드가 아닌 임의의 노드였기 때문이다.
        v3, d3, cnt3 = find_fahrest_node(v2, n, tree)
        
        #만약 이렇게 했는데도 트리의 지름을 갖는 노드 쌍이 2개 이상이라면 트리의 지름을 정답으로 반환
        if cnt3 >= 2:
            return d3
        #아니라면 가장 큰 값이 트리의 지름이고 중간값은 그보다 하나 작은 값이므로
        #트리의 지름 -1
        else:
            return d3 - 1