백준 1240 - 노드 사이의 거리(파이썬)

난이도

골드 5

 

유형

그래프 탐색

트리

해결

간단하게 (n+1)의 제곱 크기의 배열로 구현 가능하다. 더 빠르게 풀고싶다면 딕셔너리를 이용해주면 된다.

n의 크기도 그리 크지 않고 m의 크기도 그리 크지 않아서 m번의 요청이 들어올 때 마다 값을 저장하지 않고 탐색을 해 주어도 시간적으로 넉넉하다. dfs나 bfs나 다익스트라 중 자신이 편한 것을 이용해주면 된다.

거리가 10000이하의 정수 즉 0이나 음수가 들어갈 수도 있지만 구조가 트리 구조기 때문에 사이클이 형성되어 있지 않으므로 별로 상관 없다.

코드

배열 사용(느림)

from sys import stdin
from collections import deque
from math import inf
input=stdin.readline

n,m=map(int,input().split()) # 값 입력 
tree=[[inf for i in range(n+1)] for j in range(n+1)] # 트리 ((n+1) * (n+1) 사이즈)

for i in range(n-1):
    a,b,dist=map(int,input().split())
    tree[a][b]=tree[b][a]=dist

def bfs(start,end):
    queue=deque([(start,0)]) # 시작지와 누적 거리
    v=[0]*(n+1)
    while queue:
        now,totaldist=queue.popleft()
        if now==end:
            return totaldist
        for i in range(1,n+1):
            if tree[now][i]!=inf and v[i]==0:
                v[i]=1
                queue.append((i,totaldist+tree[now][i]))
                

for j in range(m):
    a,b=map(int,input().split())
    print(bfs(a,b))

딕셔너리 사용(빠름)

from sys import stdin
from collections import deque,defaultdict
from math import inf
input=stdin.readline

n,m=map(int,input().split()) # 값 입력 
tree={i:defaultdict(int) for i in range(1,n+1)}

for i in range(n-1):
    a,b,dist=map(int,input().split())
    tree[a][b]=tree[b][a]=dist

def bfs(start,end):
    queue=deque([(start,0)]) # 시작지와 누적 거리
    v=[0]*(n+1)
    while queue:
        now,totaldist=queue.popleft()
        if now==end:
            return totaldist
        for i in tree[now]:
            if v[i]==0:
                v[i]=1
                queue.append((i,totaldist+tree[now][i]))
                

for j in range(m):
    a,b=map(int,input().split())
    print(bfs(a,b))

좋은 / 비슷한 문제