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

문제

https://programmers.co.kr/learn/courses/30/lessons/68937

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

풀이

1. 트리의 지름

트리에서 어떤 두 노드를 선택해서 양쪽으로 쫙 당길 때, 가장 길게 늘어나는 경우가 있을 것이다.

이럴 때 트리의 모든 노드들은 이 두 노드를 지름의 끝 점으로 하는 원 안에 들어가게 된다.

이런 두 노드 사이의 경로의 길이를 트리의 지름이라고 한다.

https://www.acmicpc.net/problem/1967

1967번: 트리의 지름

2. 문제 풀이

• 임의의 노드(1)에서 가장 먼 노드 X를 찾는다.
• 노드 X에서 가장 먼 노드 Y를 찾는다.
  • 노드 Y가 두개 이상이면, 지름을 찾은 것이다.
• 노드 Y에서 가장 먼 노드 Z를 찾는다.
  • 노드 Z가 두개 이상이면, 지름을 찾은 것이다.
  • 노드 Z가 한개 이하라면, 가장 가까운 노드를 선택해야 하므로 지름-1을 반환한다.

[EX1]

Find x : 1에서 가장 먼 노드는 4

Find y : 4에서 가장 먼 노드는 1

Find z : 1에서 가장 먼 노드는 4

return 2

 

1, 4, 3을 선택했을 때

(1, 4사이의 거리 = 3), (2, 4 사이의 거리 = 3), (1, 3 사이의 거리 = 2) 의 중간값은 

지름인 3-1 = 2이다.

[EX2]

Find x : 1에서 가장 먼 노드는 4, 3, 2

Find y : 4에서 가장 먼 노드는 3, 2, 1

return 2

코드

from collections import defaultdict
graph = defaultdict(list)

def search(root, n) : 
    stack = [(root, 0)]
    passed = [-1 for i in range(n+1)]
    weight = defaultdict(list)
    while stack :
        node, w = stack.pop()
        passed[node] = w
        weight[w].append(node)
        
        for c in graph[node] : 
            if passed[c] == -1 : 
                stack.append((c, w+1))
    return weight

def solution(n, edges):    
    for c, p in edges : 
        graph[c].append(p)
        graph[p].append(c)
    
    # 임의의 노드 1에서 가장 멀리 있는 노드 X를 찾는다.
    weight = search(1, n)
    diameter = max(weight)
    node_X = weight[diameter][0]
    
    # 노드 X에서 가장 먼 노드 Y를 찾는다.
    weight = search(node_X, n)
    diameter = max(weight)
    if len(weight[diameter])>=2 : return diameter
    
    # 노드 Y에서 가장 먼 노드 Z를 찾는다.
    node_Y = weight[diameter][0]
    weight = search(node_Y, n)
    diameter = max(weight)
    if len(weight[diameter])>=2 : return diameter
    return diameter-1