BOJ) MooTube

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

 

 

정말 기본적인 BFS문제였다.

기초에 충실하게 bfs를 작성하고 제출을 눌렀으나 시간초과가 발생했다.

import math
from collections import deque
import sys

input = sys.stdin.readline
INF = math.inf

n, q = map(int, input().split())


arr = [[] for _ in range(n+1)]

for _ in range(n-1):
    a, b, c = map(int, input().split())

    arr[a].append([b, c])
    arr[b].append([a, c])


for _ in range(q):
    k, start = map(int, input().split())
    visit = [False] * (n + 1)
    visit[start] = True
    cnt = 0

    q = deque()
    q.append([start, INF])

    while q:
        node, value = q.popleft()

        for i in arr[node]:
            new_node, new_value = i[0], i[1]

            if new_value >= k and visit[new_node] == False:
                cnt += 1
                q.append((new_node, new_value))
                visit[new_node] = True

    print(cnt)

 

내가 제출한 코드이다.

어디서 시간초과가 발생했는지 한참을 찾다 다른 사람의 코드를 봤다.

 

from collections import defaultdict, deque

def bfs(K, num):
  visited = [0]*(len(video.keys())+1)
  cnt = 0
  q = deque()
  q.append(num)
  visited[num] = 1

  while q:
    n = q.popleft()
    for v in video[n]:
      if not visited[v[0]] and v[1] >= K:
        visited[v[0]] = 1
        cnt += 1
        q.append(v[0])

  return cnt

N, Q = map(int, input().split())
video = defaultdict(list)
for _ in range(N-1):
  p, q, r = map(int, input().split())
  video[p].append([q, r])
  video[q].append([p, r])

for _ in range(Q):
  k, v = map(int, input().split())
  print(bfs(k, v))

 

각 노드의 연결 리스트를 담는 변수를 리스트가 아닌 딕셔너리로 선언하는 곳에서만 차이가 발생할 뿐 그 외 별다른 차이가 없다.

사실 리스트보다 딕셔너리가 언제나 O(1)로 복잡도는 1이지만 리스트에서의 인덱스로 접근하는 복잡도와 비교할 만할까 하는 생각이다.