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

난이도: Silver 2
번호: 25516
생성일: March 9, 2023 2:39 PM
알고리즘 분류: 그래프 이론, 그래프 탐색, 깊이 우선 탐색, 너비 우선 탐색, 트리
언어: C++

25516번: 거리가 k이하인 트리 노드에서 사과 수확하기

C++

#include <cstdio>
#include <queue>
#include <vector>

using namespace std;

int N, K;
vector<vector<int>> map;
vector<int> apple;
vector<bool> visit;

int BFS()
{
    queue<int> Q;
    Q.push(0);
    visit[0] = true;
    int depth = 0;
    int cnt = apple[0];

    while (!Q.empty() && depth < K)
    {
        int size = Q.size();

        while (size--)
        {
            int cur = Q.front();
            Q.pop();

            for (int i : map[cur])
            {
                if (visit[i])
                    continue;

                Q.push(i);
                visit[i] = true;
                cnt += apple[i];
            }
        }
        depth++;
    }
    return cnt;
}

int main(void)
{
    scanf("%d %d", &N, &K);

    map.resize(N + 1);
    apple.resize(N + 1);
    visit.resize(N + 1);

    for (int i = 0; i < N - 1; i++)
    {
        int a, b;
        scanf("%d %d", &a, &b);
        map[a].push_back(b);
        map[b].push_back(a);
    }

    for (int i = 0; i < N; i++)
        scanf("%d", &apple[i]);

    printf("%d", BFS());

    return (0);
}