[PS / Python] 백준 #15592 MooTube(Silver)

문제 정보

링크: https://www.acmicpc.net/problem/15591
태그: 그래프, BFS

풀이

Q번 k와 v를 입력 받으면서 v를 시작점으로 하여 bfs를 돌린다.
bfs과정에서 인접 리스트를 돌며, 만약 간선의 가중치가 현재 노드에 오기까지 만난 간선의 최소값보다 작으면 그 값을 바꾼다.
한 노드에 도착했을 때마다(큐에 넣을 때마다) 그 가중치의 최소값이 k보다 크다면 cnt값을 증가시킨다.

시행착오

풀이 자체는 맞는 것 같은데, 자꾸 시간초과가 떠서 결국 다른 분의 코드를 참고했다.
나와 정답 코드의 다른 점은

• visited 배열 초기화 -> visited = [0] * (N+1)
• queue를 list를 사용하지 않고, collections의 deque를 사용 (q.pop(), q.append())
• 나는 인접행렬을 사용했지만, 정답코드에서는 인접 리스트 사용 -> adj = [[] for _ in range(N+1)], adj[p].append((q, u))
• input = sys.stdin.readline으로 바꾸어 사용

아마도 나는 인접 리스트의 5000개를 다 돌기 때문에 시간초과가 떴지만, 인접 리스트를 사용하여 실행 시간을 줄인 것 같다.
또한 input = sys.stdin.readline으로 바꾼 것도 영향이 있는 것 같다.

import sys
from collections import deque

input = sys.stdin.readline

def bfs(v: int, k: int, adj: list, N: int):
    cnt = 0
    visited = [0] * (N+1)
    visited[v] = 1
    queue = deque()
    queue.append((v, 1e9))
    while queue:
        cur, u = queue.pop()
        for next, nu in adj[cur]:
            if visited[next]:
                continue
            visited[next] = 1
            nu = min(u, nu)
            if nu >= k:
                cnt += 1
            queue.append((next, nu))
    return cnt
    

if __name__ == "__main__":
    N, Q = map(int, input().split())
    adj = [[] for _ in range(N+1)] # 인접리스트

    for _ in range(N-1):
        p, q, u = map(int, input().rstrip().split())
        adj[p].append((q, u))
        adj[q].append((p, u))
    
    for i in range(Q):
        k, v = map(int, input().rstrip().split())
        ans = bfs(v, k, adj, N)
        print(ans)