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

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

코드

#include <iostream>
#include <vector>
using namespace std;

#define FASTIO ios::sync_with_stdio(false);cin.tie(NULL);
#define MAX 100001

int n, k, ans;
vector<int> v[MAX];
int cnt[MAX];

void dfs(int now, int ncnt) {
	if (ncnt > k) return;
	ans += cnt[now];
	for (int i = 0; i < v[now].size(); i++)
		dfs(v[now][i], ncnt + 1);
}

int main() {
	FASTIO;
	cin >> n >> k;

	for (int i = 0, s, e; i < n - 1; i++) {
		cin >> s >> e;
		v[s].push_back(e);
	}
	
	for (int i = 0; i < n; i++) cin >> cnt[i];
	dfs(0, 0);
	cout << ans;
	return 0;
}

설명

트리를 생성한 후에 루트에서부터 DFS를 돌려서 K번까지 자식 노드로 이동합니다. 이 과정에서 노드별로 가진 사과의 개수를 더해나가고 마지막에 사과의 개수를 출력하면 됩니다.

 

느낀 점

DFS를 사용해서 한 번에 풀었습니다~