거리가 k이하인 트리 노드에서 사과 수확하기[백준 25516]

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

from collections import deque

n, k = map(int, input().split())
gra = [[] for _ in range(n)]
for _ in range(n-1):
    a, b = map(int, input().split())
    gra[a].append(b)
    gra[b].append(a)

isapple = list(map(int, input().split()))
visit = [0] * n

def bfs(start):
    q = deque()
    q.append((start, 0))  # (현재 노드, 깊이)
    visit[start] = 1
    ans = 0  # 사과 개수
    if isapple[start] == 1:
        ans += 1

    while q:
        current, depth = q.popleft()
        if depth >= k:  # 깊이가 k를 넘으면 종료
            continue

        for next_node in gra[current]:
            if visit[next_node] == 0:
                visit[next_node] = 1
                q.append((next_node, depth + 1))
                if isapple[next_node] == 1:
                    ans += 1

    print(ans)  # 결과 출력

bfs(0)