[baekjoon 16437] 양 구출 작전 (DFS, 트리) (C++)

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

16437번: 양 구출 작전

각 섬에 있는 양들을 1번 섬으로 옮겨야 한다. 하지만 경로 상에 늑대가 있다면 양은 늑대의 수만큼 줄어든다.

 

각 섬에서 1번 섬으로 가는 루트는 유일하므로 트리 구조를 생각할 수 있고,

늑대는 양을 최대 1마리만 잡아먹으므로 섬에 계속 남아서 오는 양을 잡아먹는 것이 아닌, 한 마리 잡아먹으면 늑대 역시 줄여줘야 한다.

 

처음엔 DFS로 모든 양의 위치에서 루트까지 탐색을 해주었더니 시간 초과가 나왔고(이 방법은 n^2이었다.) 따라서 트리 구조를 이용해서 트리를 한번 탐색해서 결과를 도출해야겠다고 생각을 했다.

 

따라서 여러 시행착오 끝에, 1번 섬부터 탐색을 했고 내려갈 때 값을 갖고 가는 방법을 썼다가 애를 먹었다.

 

결론적으로

1. 1번 섬부터 탐색하되, 리프 노드까지 가서 양의 수를 가져오는 방법으로 풀었다.

2. 리프 노드에서 섬에 양이 있다면 리턴 값으로 가져오고, 늑대가 있다면 무시한다.

(양이 가는 경로에 늑대가 있어야 늑대를 사용하기 때문)

3. 이후부터는 dfs를 실행한 노드의 늑대 or 양의 상태로 더해주고, 빼주며 루트까지 재귀적으로 다시 오면 된다.

 

+ 결괏값은 int의 범위를 넘어간다.

#include <iostream>
#include <vector>
#include <algorithm>

#define MAX 123457
using namespace std;

typedef long long ll;
int N, a, p, wolf[MAX], island[MAX];
ll answer;
vector<int> v[MAX];
char animal;

long long DFS(int startX) {
	int vSize = v[startX].size();

	if (vSize == 0) {
		if (wolf[startX])
			return 0; // 리프노드가 늑대이면 무시
		else
			return island[startX]; // 양이면 갖고 올라간다
	}
	long long tempResult = 0;

	for (int i = 0; i < vSize; i++) {
		int nx = v[startX][i];
		tempResult += DFS(nx);
	}

	// 밑의 노드들 계산한 값 + 현재 노드의 상태를 추가
	if (wolf[startX]) {
		tempResult -= island[startX];
	}
	else
		tempResult += island[startX];

	if (tempResult < 0) // 만약 이 노드의 값이 0이면 무시
		return 0;
	else
		return tempResult;
}

int main() {
	ios::sync_with_stdio(false);
	cin.tie(0);
	cin >> N;
	for (int i = 2; i <= N; i++) {
		cin >> animal >> a >> p;
		if (animal == 'W')
			wolf[i] = true; // 늑대여부
		island[i] = a; // 값
		v[p].push_back(i); // 트리 연결
	}

	cout << DFS(1);
}