취미로 코딩하기

※ 글쓴이는 취미로 코딩을 익혀보는 사람이라 정확하지 않은 내용을 담고 있을 수 있다 ※

이번에 볼 문제는 백준 1623번 문제인 신년 파티이다.
문제는 아래 링크를 확인하자.

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

최적해를 가질 때의 그 구성 노드를 역추적한 결과를 요구하는 DP문제이다.

노드i를 사용하지 않을 때 노드i를 루트로 갖는 서브트리의 "날라리 분위기"의 최댓값을 dp0[i], 노드i를 사용할 때 노드i를 루트로 갖는 서브트리의 "날라리 분위기"의 최댓값을 dp1[i]로 표현하면, 전체 트리의 리프노드서부터 루트방향으로 위상정렬순으로 모든 노드를 순회하며 dp0[i]와 dp1[i]값을 채워나가는 것으로 최적해를 구해낼 수 있다.

이 때, 어떤 dp0[j] 또는 dp1[j]가  dp0[i]의 값에 영향을 주었는지와 dp1[i]의 값에 영향을 주었는지를 인접리스트의 형태로 보관해, 해당 인접리스트를 이용하여 답을 역추적해낼 수 있다.

문제에서 주어지는 "날라리 기질"은 음수값이 주어질 수 있음에 유의하여 구현하자.

아래는 제출한 소스코드이다.

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

int parent[200001];
int indegree[200001];
int dp0[200001], dp1[200001];
vector<pair<int,int>> dp0node[200001], dp1node[200001]; // {자식노드, 0/1}

vector<int> ans;

void backtrack1(int cur);

void backtrack0(int cur) {
	for (auto node : dp0node[cur]) {
		if (node.second) backtrack1(node.first);
		else backtrack0(node.first);
	}
}

void backtrack1(int cur) {
	ans.emplace_back(cur);
	for (auto node : dp1node[cur]) {
		backtrack0(node.first);
	}
}

int main() {
	ios::sync_with_stdio(0);
	cin.tie(0);

	int N; cin >> N;
	for (int i = 1; i <= N; i++) cin >> dp1[i];
	for (int i = 2; i <= N; i++) {
		int& tmp = parent[i];
		cin >> tmp;
		indegree[tmp]++;
	}

	queue<int> que;
	for (int i = 1; i <= N; i++) {
		if (indegree[i] == 0) que.push(i);
	}

	while (!que.empty()) {
		int cur = que.front(); que.pop();
		int par = parent[cur];
		
		if (dp0[cur] >= dp1[cur] && dp0[cur] > 0) {
			dp0[par] += dp0[cur];
			dp0node[par].emplace_back(make_pair(cur, 0));
		}
		else if (dp1[cur] > dp0[cur] && dp1[cur] > 0) {
			dp0[par] += dp1[cur];
			dp0node[par].emplace_back(make_pair(cur, 1));
		}

		if (dp0[cur] > 0) {
			dp1[par] += dp0[cur];
			dp1node[par].emplace_back(make_pair(cur, 0));
		}

		indegree[par]--;
		if (indegree[par] == 0) que.push(par);
	}

	cout << dp1[1] << ' ' << dp0[1] << '\n';

	backtrack1(1);
	sort(ans.begin(), ans.end());
	for (auto x : ans) cout << x << ' ';
	cout << -1 << '\n';
	
	ans.clear();
	backtrack0(1);
	sort(ans.begin(), ans.end());
	for (auto x : ans) cout << x << ' ';
	cout << -1 << '\n';
}