백준 1623번 - 신년 파티 (C++)

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

1623번: 신년 파티

 

트리에서 다이나믹 프로그래밍인 것은 바로 파악할 수 있었다. 

그리고 날라리 점수를 파악하는 것 까지는 성공했으나 참가하는 목록을 확정하는 과정이 쉽지 않았다. 먼저 트리를 그려 파악하면 다음과 같다.

 

 

 

트리에서의 다이나믹 프로그래밍은 dfs를 돌면서 첫 방문에 필요 정보를 초기화하고 이후 dfs가 끝났을 때 정보들을 초기화하게 된다. 

 

코드로 적어보면

 

void dfs(int n) {
	dp[n][0] = 0;
	dp[n][1] = score[n];

	for (int i = 0; i < node[n].size(); i++) {
		int next_n = node[n][i];
		dfs(next_n);
	}
}

 

이런식으로 먼저 첫 방문에는 dp값을 초기화 시킨다. 여기서 dp[n][0]은 참가하지 않는 것이고 dp[n][1]은 참가한다는 뜻이다. 그리고 dfs 를 계속 돈다. 이제 dfs가 끝에 다다랐을 경우 위의 사진처럼 리프노드인 4번은 참가시엔 15, 참가안할시엔 0이된다. 그리고 2번으로 다시 돌아갔을 때 

 

dp[2][0] => 부모가 참가하지 않으므로 자식이 참가해도되고 안해도 된다.

dp[2][1] => 내가 무조건 참가한다. 자식은 참가하지 못한다.

 

이것이 이루어진다. 이것을 코드로 바꿔보면

 

void dfs(int n) {
	dp[n][0] = 0;
	dp[n][1] = score[n];

	for (int i = 0; i < node[n].size(); i++) {
		int next_n = node[n][i];
		dfs(next_n);
		dp[n][1] = dp[next_n][0] + dp[n][1];
		dp[n][0] = max(dp[next_n][0], dp[next_n][1]);
	}
}

 

여기서 끝났으면 이 문제는 골드1 난이도가 아닐 것이다. 이 문제는 더 나아가서 지나온 로그를 묻는다. 게다가 사장이 참가했을 때, 안했을 때를 둘 다 물어보고 있다. (참으로 가혹하다)

그래서 이것들을 추적하기 위한 단계가 필요하다. dfs를 방문하면서 값은 계속 갱신되기 때문에 1번의 dfs로는 문제를 해결할 수 없다. 그래서 표시를 해두고 로그를 다시 쫓는 단계가 추가적으로 필요하다. 때문에 check 배열을 선언한다.

 

void dfs(int n) {
	dp[n][0] = 0;
	dp[n][1] = score[n];

	for (int i = 0; i < node[n].size(); i++) {
		int next_n = node[n][i];
		dfs(next_n);
		check[n] = true;
		dp[n][1] = dp[next_n][0] + dp[n][1];
		if (dp[next_n][0] > dp[next_n][1]) {
			dp[n][0] += dp[next_n][0];
			check[next_n] = false;
		}
		else {
			dp[n][0] += dp[next_n][1];
			check[next_n] = true;
		}
	}
}

 

check 배열은 먼저 부모 노드는 참가한다고 우선 true로 표시한다. 그리고 dp[n][0] 값을 결정할 때, 만약 자식이 참가하도록설정된다면 자식을 true로, 자식이 불참으로 설정된다면 false로 설정한다.

 

그럼 이걸 나중에 어떻게 쓸 수 있을까. log_dfs를 선언해 1번부터 dfs를 돌며, 사장이 참가했을 때 참가하는 인원들과 사장이 불참시 참가하는 인원들을 알아낼 수 있다.

왜냐하면 위 dfs의 종착지는 언제나 1번노드이기 때문에 dp[1][0]과 dp[1][1]을 결정할 때 true false도 그와 맞춰 갱신되기 때문이다.

 

void log_dfs(int n, bool ch, int idx) {
	if (ch) {
		if (idx == 1)
			log_1.push_back(n);
		else
			log_0.push_back(n);
	}

	for (int i = 0; i < node[n].size(); i++) {
		int next = node[n][i];
		if (ch) {
			log_dfs(next, false, idx);
		}
		else {
			if (check[next]) log_dfs(next, true, idx);
			else log_dfs(next, false, idx);
		}
	}
}

 

먼저, log_dfs(1,true, 1) 이라하면

만약 현재 check 가 true라면 log 벡터에 넣는다. 

그리고 다음 노드를 확인한다. 만약 현재 노드가 true였다면 다음 노드는 언제나 false일 수 밖에 없다. 그러나 만약 이 노드가 false였다면 다음 노드의 check배열을 확인해야한다.

 

위를 통해 트리 사진의 

 

     2

4       5

 

관계를 해결할 수 있다. 여기서 사장이 참여시 정답은 1번과 4번이다. 2번 노드는 참여하지 않는다. 때문에 4번이 참여하는가 하지 않는가를 확인하고 삽입하고, 5번도 참여하는가 하지 않는가를 확인하고 삽입한다.

 

트리에서 다이나믹 트리도 까다로운 문제인데 여기서 추적까지 필요한 매우 어려운 문제였다.

 

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

map<int, int> score;
vector<int> node[200001];
bool check[200001] = { false, };

vector<int> log_0;
vector<int> log_1;

int dp[200001][2];
void dfs(int n) {
	dp[n][0] = 0;
	dp[n][1] = score[n];

	for (int i = 0; i < node[n].size(); i++) {
		int next_n = node[n][i];
		dfs(next_n);
		check[n] = true;
		dp[n][1] = dp[next_n][0] + dp[n][1];
		if (dp[next_n][0] > dp[next_n][1]) {
			dp[n][0] += dp[next_n][0];
			check[next_n] = false;
		}
		else {
			dp[n][0] += dp[next_n][1];
			check[next_n] = true;
		}
	}
}

void log_dfs(int n, bool ch, int idx) {
	if (ch) {
		if (idx == 1)
			log_1.push_back(n);
		else
			log_0.push_back(n);
	}

	for (int i = 0; i < node[n].size(); i++) {
		int next = node[n][i];
		if (ch) {
			log_dfs(next, false, idx);
		}
		else {
			if (check[next]) log_dfs(next, true, idx);
			else log_dfs(next, false, idx);
		}
	}
}

int main() {
    ios_base::sync_with_stdio(false);
	cin.tie(NULL);
	int N;
	cin >> N;

	for (int i = 1; i <= N; i++) {
		int tmp;
		cin >> tmp;
		score.insert({ i, tmp });
	}

	for (int i = 2; i <= N; i++) {
		int tmp;
		cin >> tmp;
		node[tmp].push_back(i);
	}

	dfs(1);
	log_dfs(1, true, 1);
	log_dfs(1, false, 0);
	cout << dp[1][1] << " " << dp[1][0] << '\n';
	sort(log_0.begin(), log_0.end());
	sort(log_1.begin(), log_1.end());
	log_0.push_back(-1);
	log_1.push_back(-1);
	for (auto a : log_1) cout << a << " ";
	cout << '\n';
	for (auto a : log_0) cout << a << " ";
	cout << '\n';
}