파이썬으로 풀어보는 백준 1967번: 트리의 지름

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

1967번: 트리의 지름

 

 

내 풀이:

from collections import deque
 
class Node:
    def __init__(self, num, val, parent = None):
        self.num = num
        self.val = val
        self.parent = parent
        self.children = []
 
 
class Tree:
    def __init__(self, N):
        self.root = Node(1, 0)
        self.current = None
        self.node_list = [0] * (N+1)
        self.node_list[1] = self.root
 
 
    def append(self, parent, num, val):
        now = self.node_list[parent]
        new = Node(num, val, now)
        now.children.append(new)
        self.node_list[num] = new
 
 
    def BFS(self, node, N):
        q = deque()
        q.append((node, 0))
        visited = [False] * (N+1)
        visited[node.num] = True
        max_length = 0
        ans = 0
        while q:
            data = q.popleft()
            self.current = data[0]
            if data[1] > max_length:
                max_length = data[1]
                ans = data[0]
 
            if self.current.parent and not visited[self.current.parent.num]:
                visited[self.current.parent.num] = True
                q.append((self.current.parent, data[1] + self.current.val))
 
            for child in self.current.children:
                if not visited[child.num]:
                    visited[child.num] = True
                    q.append((child, data[1] + child.val))
 
        self.current = ans
        return max_length
 
 
N = int(input())
tree = Tree(N)
 
for i in range(N-1):
    parent, child, val = map(int,input().split())
    tree.append(parent, child, val)
 
tree.BFS(tree.root, N)
print(tree.BFS(tree.current, N))

Python3, 540ms

먼저 클래스를 이용하여 트리를 구현한 경우이다. 문제를 풀기 위한 알고리즘으로는 트리의 Root에서 거리가 가장 먼 노드를 찾아 저장한 후, 저장한 노드를 출발 지점으로 하여 가장 먼 노드까지의 거리를 출력하였다. 이 방법의 단점은 트리를 만들기 위한 행위를 제외하고 트리 전체 탐색만 2번을 해야 한다는 점이다. 이 문제는 최대 만 개의 노드 정보를 받아야 하기 때문에 import sys 후 input = sys.stdin.readline() 처리를 해준다면 120ms 정도의 시간이 걸리는 것을 확인할 수 있다.

 

import sys
from collections import deque
 
N = int(sys.stdin.readline())
Tree = [[0, [], 0] for _ in range(N+1)] # [[parent, child, val], ... , ]
 
for _ in range(N-1):
    parent, child, val = map(int,sys.stdin.readline().split())
    Tree[parent][1].append(child)
    Tree[child][0] = parent
    Tree[child][2] = val
 
q = deque()
q.append((1, 0))
visited = [False] * (N+1)
visited[1] = True
 
distance = 0
furthest = 0
 
while q:
    data = q.popleft()
    cur = data[0]
    res = data[1]
 
    if distance < res:
        distance = res
        furthest = cur
 
    parent = Tree[cur][0]
    children = Tree[cur][1]
    if parent and not visited[parent]:
        visited[parent] = True
        q.append((parent, res + cur[2]))
 
    for child in children:
        if not visited[child]:
            visited[child] = True
            q.append((child, res + Tree[child][2]))
 
distance = 0
q = deque()
q.append((furthest, 0))
visited = [False] * (N+1)
visited[furthest] = True
while q:
    data = q.popleft()
    cur = data[0]
    res = data[1]
 
    if distance < res:
        distance = res
        furthest = cur
 
    parent = Tree[cur][0]
    children = Tree[cur][1]
    if parent and not visited[parent]:
        visited[parent] = True
        q.append((parent, res + Tree[cur][2]))
 
    for child in children:
        if not visited[child]:
            visited[child] = True
            q.append((child, res + Tree[child][2]))
 
print(distance)

Python3, 120ms

이 풀이도 알고리즘은 동일하기 때문에 시간은 524ms가 나왔으나 readline()을 이용하여 120ms가 나온 것을 확인할 수 있다. 클래스 대신 리스트를 이용하여 트리를 구현해주기 위해 2차원 리스트를 이용해주었으며 구조는 주석으로 달아놓은 것과 같다. 클래스에 비해 구조를 머릿속으로 그리기 어렵지만 코드 길이 자체는 더 짧기 때문에 트리에 충분히 적응하여 리스트를 이용하는 것이 바람직해 보인다.