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

문제

---

풀이 과정

1) 문제 이해 및 접근

세개의 노드를 선택하여야 하므로, 모듈을 이용해 조합의 경우의 수를 구하는 방식으로 접근했습니다.

또한 모든 점에 대해서 어떤 점을 향하는 거리의 값이 필요하여 다익스트라를 이용하는 방식으로 생각했습니다.

 

그러나 노드의 개수가 너무 많아, 너무 많은 다익스트라 함수를 실행하여 시간 초과가 발생하였습니다.

이 문제를 해결하기 위해 고민하던 중 트리의 지름에 관한 문제를 접하고, 활용하게 되었습니다.

 

세개의 노드의 길이 중 최댓값을 구하면 되었으므로 트리의 지름을 구하고 트리의 지름 -1 의 길이를 구하는 방식으로 접근하였습니다.

2) 알고리즘

• 다익스트라

3) 풀이 코드

사용 언어 - Python

(1) 다익스트라, 조합 이용한 접근 - 시간 초과

import heapq
from itertools import combinations


def dijkstra(start, distance, linked, n):
    visited = [0] * (n+1)
    heap = [(0, start)]

    while heap:
        node = heapq.heappop(heap)
        if not visited[node[1]]:
            visited[node[1]] = 1
            distance[start][node[1]] = node[0]
            for e in linked[node[1]]:
                heapq.heappush(heap, (1+node[0], e))



def solution(n, edges):
    distance = [[0] * (n+1) for _ in range(n+1)]
    linked = [[] for _ in range(n+1)]
    answer = []

    for edge in edges:
        linked[edge[0]].append(edge[1])
        linked[edge[1]].append(edge[0])
    

    for i in range(1, n+1):
        dijkstra(i, distance, linked, n)

    lst_comb = combinations(range(1, n+1), 3)
    
    for e in lst_comb:
        tmp = [0, 0, 0]
        tmp[0] = distance[e[0]][e[1]]
        tmp[1] = distance[e[0]][e[2]]
        tmp[2] = distance[e[1]][e[2]]

        answer.append(sum(tmp) - max(tmp) - min(tmp))
        
    return max(answer)

 

(2) 다익스트라, 트리의 지름 이용한 접근 - 최종 코드

import heapq

def dijkstra(start, linked, n):                                     # 다익스트라, 시작 노드 / 연결 관계 리스트 / 노드의 개수
    visited = [0] * (n+1)                                           # 노드 방문 여부 리스트
    distance = [0] * (n+1)                                          # 시작 노드로부터 각 노드까지의 거리 리스트
    heap = [(0, start)]                                             # 최소힙 (거리, 시작 노드)

    while heap:                                                     
        node = heapq.heappop(heap)                                  # 최소힙에서 다음 이동 가능한 노드 정보 꺼냄
        if not visited[node[1]]:                                    # 아직 방문하지 않은 곳 일때
            visited[node[1]] = 1                                    # 방문 처리
            distance[node[1]] = node[0]                             # 해당 노드까지 거리 입력
            for e in linked[node[1]]:                               # 현재 노드와 연결된 다음 노드의 정보들
                heapq.heappush(heap, (1 + distance[node[1]], e))    # 거리를 1 추가하고 힙에 추가
    
    return distance                                                 # 노드의 거리 정보가 담긴 리스트 반환


def solution(n, edges):
    linked = [[] for _ in range(n+1)]                               # 연결 관계 리스트

    for edge in edges:                                              # 노드간 양방향 간선
        linked[edge[0]].append(edge[1])             
        linked[edge[1]].append(edge[0])
    
    distance = dijkstra(1, linked, n)                               # 최초 1회 임의의 노드에서 출발하여 가장 멀리 있는 노드를 구함
    distance = dijkstra(distance.index(max(distance)), linked, n)   # 가장 멀리 있는 노드에서 시작하여 노드의 거리를 구함

    if distance.count(max(distance)) >= 2:                          # 가장 멀리 있는 노드가 두개 이상인 경우
        return max(distance)                                        # 트리의 지름이 중간값이 되므로 트리의 지름 반환

    distance = dijkstra(distance.index(max(distance)), linked, n)   # 가장 멀리 있는 노드가 하나인 경우, 확인을 위해 가장 먼 노드에서 다시 다익스트라 실행

    if distance.count(max(distance)) >= 2:                          # 가장 멀리 있는 노드가 두개라면 트리의 지름 반환
        return max(distance)
    else:                                                           # 가장 멀리 있는 노드가 한개라면, 트리의 지름 -1 반환
        return max(distance) - 1

---

결과 및 학습한 내용

1) 어려웠던 내용

문제를 읽은 후 트리의 지름을 이용하는 문제라는 것을 쉽게 판단하지 못하였습니다.

트리의 지름을 이용하는것을 파악한 후에도 트리의 지름이 1개만 나오는 경우, 한번 더 탐색을 진행해야 하는 이유를 이해하기 어려웠습니다.

(1) 참고 사이트

• 블로그
  트리의 지름이 한개 나오는 경우, 한번 더 탐색을 진행하는 테스트 케이스를 확인할 수 있었습니다.

2) 새롭게 학습한 내용

트리의 지름의 특징을 배우게 되었습니다.

• 트리의 지름을 구하기 위해서는 최초 1회를 임의의 노드에서 탐색한 후 가장 멀리 있는 노드에서 다시 탐색하여 나온 가장 멀리 있는 노드까지의 길이가 트리의 지름이 된다.
• 트리의 지름의 개수를 구하는 경우, 최초 1회 이외에도 2번 정도 탐색을 해야할수도 있다.

---

문제 링크

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

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

 

 

※ 오류 및 오타, 다른 의견이 있는 경우 댓글을 남겨주시면 감사하겠습니다