백준 1967번: 트리의 지름

✔ 풀이를 위한 아이디어

• 그래프 이론 + 그래프 탐색 (DFS or BFS)
• 그래프로 표현되는 트리 (Tree) 의 이해
  
  

✔ 오답 코드 [틀렸습니다]

import sys

n = int(sys.stdin.readline())
tree = [[] for _ in range(n+1)]  # index 0은 사용 X
length = [0]*(n+1)
node_max = [0]*(n+1)

for _ in range(n-1):
    parent, child, value = map(int, sys.stdin.readline().split())
    tree[parent].append((child, value))  # index 0은 자녀이름, 1은 가중치

max_len = 0


def postOrder(cur_root):
    tmp = tree[cur_root]
    if len(tmp) >= 1:
        length[tmp[0][0]] = length[cur_root] + tmp[0][1]
        postOrder(tmp[0][0])
    if len(tmp) >= 2:
        length[tmp[1][0]] = length[cur_root] + tmp[1][1]
        postOrder(tmp[1][0])

    if len(tmp) == 0:
        node_max[cur_root] = length[cur_root]
    elif len(tmp) == 1:
        node_max[cur_root] = node_max[tmp[0][0]]
    elif len(tmp) == 2:
        node_max[cur_root] = max(node_max[tmp[0][0]], node_max[tmp[1][0]])
        global max_len
        if node_max[tmp[0][0]] + node_max[tmp[1][0]] - (2 * length[cur_root]) > max_len:
            max_len = node_max[tmp[0][0]] + node_max[tmp[1][0]] - (2 * length[cur_root])


postOrder(1)
print(max_len)

접근 방식

• 문제의 힌트에서는 DFS를 활용해 풀으라고 말하고 있었다. 하지만, 나는 DFS 보다 후위 순회를 하는 것이 문제에서 요구하는 답을 구하는 데 더욱 효율적일 것이라고 생각했기 때문에 재귀 호출을 활용한 후위 순회로 코드를 짜보았다.

알고리즘 설명

• 이 알고리즘은, 기본적으로 후위 순회를 하며 Root로부터 현재 노드까지의 거리를 length 배열에 저장한다.
• 또한, 각 노드의 자식 노드들의 length 값을 비교하여 그 중 최댓값을 node_max 배열에 저장한다. 이는 후위 순회의 흐름을 해치지 않고, (다시 자식 노드 방향으로 내려가며 탐색하지 않고) 부모 노드에서 즉시 최댓값을 구하기 위함이다. 자식 노드가 없을 경우에는 부모 노드의 length 값을 저장한다.
• 자식 노드가 2개인 부모 노드를 만날 경우, 자식 노드들의 node_max 값과 부모 노드의 length 값을 활용해 최댓값을 갱신한다. (자식노드의 length 값에서 부모 노드의 length 값을 빼면 자식노드부터 부모노드까지의 구간 거리를 구할 수 있다.)

한계점

• 그러나 이 코드는 자식 노드가 3개 이상일 경우를 고려하지 않고 짠 코드이다. (문제에서 이진트리라는 조건은 없었다!)
• 또한 자식 노드가 2개인 경우에 대해서만 트리의 지름을 계산하므로, 편향 트리의 경우 잘못된 답을 도출하게 된다.
• 자식 노드가 3개 이상인 경우는 어떻게든 처리할 수 있었지만, 후위 순회의 방식으로는 편향 트리의 경우를 처리할 수가 없었다. 따라서 풀이 방식을 다시 생각해내야 했다.

✔ 정답 코드 [컴파일 에러 -> 틀렸습니다 -> 해결]

import sys
from collections import deque

n = int(sys.stdin.readline())
tree = [[] for _ in range(n+1)]  # index 0은 사용 X

for _ in range(n-1):
    parent, child, value = map(int, sys.stdin.readline().split())
    tree[parent].append((child, value))  # index 0은 자녀이름, 1은 가중치
    tree[child].append((parent, value))  # 양방향으로 해야 아래서부터 탐색해서 올라가는 2번째 BFS 가능 (틀렸습니다)

length = [0]*(n+1)
visited = [0]*(n+1)  # 이거 안 쓰면 한방향으로만 탐색 못함
max_len = 0
max_node = -1


def BFS(root):
    queue = deque([root])
    visited[root] = 1
    while queue:
        cur_root = queue.popleft()
        for i in range(len(tree[cur_root])):
            if visited[tree[cur_root][i][0]] == 1:
                continue
            length[tree[cur_root][i][0]] = length[cur_root] + tree[cur_root][i][1]
            queue.append(tree[cur_root][i][0])
            visited[tree[cur_root][i][0]] = 1
            global max_len, max_node  # if 문에 넣기 전에 미리 global 변수 선언하기 (컴파일 에러)
            if length[tree[cur_root][i][0]] > max_len:
                max_len = length[tree[cur_root][i][0]]
                max_node = tree[cur_root][i][0]


BFS(1)
length = [0]*(n+1)
visited = [0]*(n+1)
max_len = 0

BFS(max_node)
print(max_len)

수정된 접근 방식

• 트리의 지름을 이루는 양 끝점은 '필연적으로' 말단 노드 (leaf node) 가 될 수밖에 없다. 따라서 DFS나 BFS를 돌며 Root 노드로부터 가장 최대의 거리에 있는 말단 노드를 찾는다.
• 찾은 말단 노드로부터 역으로 탐색을 하며 가장 멀리 떨어져 있는 노드를 찾는다.
• 즉, 이 문제는 DFS나 BFS를 두 번 쓰면 간단히 풀리는 문제인 것이다. Root 노드로부터 한번, 찾은 말단 노드로부터 한번이다. 두 번째 탐색에서 찾아낸 최대 거리가 즉 트리의 지름이 될 것이다. (나는 BFS를 사용하였다.)

코드 설명

• 우선, 이전 풀이와 다르게 child 노드에도 parent에 대한 정보를 넣어줘야 한다. 양방향 그래프로 구성해야 첫 번째 탐색에서 찾아낸 최대 거리의 말단 노드로부터 역으로 탐색할 수 있기 때문이다.
• 또한, visited 배열을 반드시 활용해야 한다. 사실 visited 배열을 쓰기 싫어서 BFS를 선택한 것이긴 한데, 결국 쓰게 되었다. 양방향으로 구성되어 있기 때문에 만약 안 쓴다면, 일정한 방향으로 쭉 탐색해나갈 수 없을 것이기 때문이다.
• BFS 함수 자체는 매우 익숙하다. 다만, length 배열에 root 노드로부터 현재 탐색 중인 노드의 자식노드들까지의 거리를 저장해나가며 탐색한다.
• 이에 더해, max_len과 max_node 값을 끊임없이 갱신해나가며 탐색이 끝났을 때에는 가장 멀리 떨어진 노드를 찾을 수 있도록 한다.
• 두 번째 BFS에서는 첫 번째 BFS에서 찾은 max_node로부터 탐색을 시작하되, 탐색 전에 초기화 해줘야할 값들을 꼼꼼히 초기화 해준다.

이 문제는 복잡한 프로그래밍 스킬보다, 문제를 단순화 시키는 능력을 요구하는듯 하다. 이런 능력은 어떻게 하면 기를 수 있을까?