백준 1949 우수 마을

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

1949번: 우수 마을

 

Tree Dp 문제

 

돌이켜생각해보면 tree는 참 재밌는 구석이 많은 자료구조같다. 지름 구하는 것도 참 신기하고...

 

양방향 간선이므로 어떤 노드를 root로 선택해도 된다. 그래서 생각하기 편하도록 1번 노드를 Root 노드로 잡았다. 

 

어떤 분은 트리 순회하는 순서를 처음부터 vector에 저장하시던데, 내 코드는 그러진 않았다.

 

코드를 보면 이해가 더 잘되실 것 같다.

 

/* 우수 마을 */

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

using namespace std;

void input();
void solve();

int N, weight[10000 + 1], Ans;
int treeDp[10000 + 1][2];

vector<int> house[10001];

int main(void) {
	ios_base::sync_with_stdio(false);
	cin.tie(NULL); cout.tie(NULL);
	
	input();
	solve();

	return 0;
}

void input() {
	cin >> N;
	for (int i = 1; i <= N; i++)
		cin >> weight[i];

	for (int i = 1; i < N; i++) {
		int a, b; cin >> a >> b;

		house[a].push_back(b);
		house[b].push_back(a);
	}
}

int recur( int curr, bool include, int before) {
	int& ret = treeDp[curr][include];

	// 중복은 뺀다. 
	if (ret != 0) {
		return ret;
	}

	int summation = 0;
	// 현재 마을은 우수 마을이다.
	if (include) {
		summation += weight[curr];

	/* 현재 마을이 우수 마을이면 다음 마을은 무조건 우수마을이 아니다. */
		for (auto nxt_node : house[curr]) {
			if (nxt_node == before) continue;
			summation += recur( nxt_node, false, curr);
		}
	}
	
	// 현재 마을은 우수 마을이 아니다.
	else {
		for (auto nxt_node : house[curr]) {
			if (nxt_node == before) continue;

	/* 현재 마을이 우수 마을이 아닐 때, 다음 마을은 우수 마을일 수도 있고 아닐 수도 있다.*/
	/* 만약 우수 마을이 아닌 경우가 세 번 이상이면 어떡하는가, 질문이 나올 수도 있지만 
	   하나의 마을이라도 더 포함하는 게 주민 수를 늘리기 때문에 세 번 연속 우수마을이 아닌 경우보다
	   적어도 하나의 마을이 우수마을이 더 많은 주민을 포함하고, 정답이다.
	   그러므로 세 번 연속 우수마을이 아닌 경우를 따로 처리하는 로직을 넣지 않았다. 
	*/
			summation += max(recur(nxt_node, true, curr), recur(nxt_node, false, curr));
		}
	}

	return ret = summation;
}

 void solve() { 
	 // 첫번째 마을이 우수마을인 경우 vs 우수마을이 아닌 경우
	 Ans = max(recur(1, true, -1), recur(1, false, -1));

	 cout << Ans << '\n';
}