[백준] 17831 - 대기업 승범이네

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

17831번: 대기업 승범이네

 

tree DP 문제이다. 어떤 부모-자식 관계의 노드 2개의 쌍들을 만들어서 그 쌍들의 곱의 합의 최댓값을 구하는 문제이다.

본인은 dp 배열을 다음과 같이 정의하였다.

dp[i][j] = i번째 노드가 (j == 1이면 멘토, j == 0이면 멘토가 아닐 때), i 번째 노드를 포함한 서브트리의 시너지의 최댓값

따라서 j가 1일때와 0일 때를 나누어서 생각해 보자.

j == 1

j가 1이면 해당 노드를 멘토로 정한다는 의미이다. 따라서 해당 노드의 자식 노드 중 하나는 무조건 멘티가 되어야 한다. 이 부분을 dp로 모든 경우를 구하면 된다.

현재 노드 node의 자식 하나하나를 child라 하고, 현재 노드의 멘티로 선택한 자식 노드를 x라 하자.

그럼 dp[node][1]은 다음과 같이 될 것이다.

dp[node][1] = max(dp[node][1], work[node] * work[x] + x_제외한_모든_자식_노드에_대해서_sum(max(dp[child][0], dp[child][1]) + dp[x][0])

이 부분을 완전히 for문으로 돌리게 되면 특수한 케이스에서 시간 초과가 날 것이다.

 

그런데 x_제외한_모든_자식_노드에_대해서_sum(max(dp[child][0], dp[child][1])) 이 부분을 잘 생각해 보자.

모든 자식 노드에 대해서 sum의 결과는 바로 dp[node][0]이다. 즉, 현재 자신을 멘토로 선택하지 않은 값과 같다.

그래서 위 식은 다음과 같이 쓸 수 있다.

dp[node][1] = max(dp[node][1], work[node] * work[x] + dp[node][0] - max(dp[x][0], dp[x][1]) + dp[x][0])

j == 0

j가 0이면 해당 노드의 자식 노드와는 연관이 없게 되므로, 모든 자식 노드에 대해서 dp[i][j] += max(dp[child][0], dp[child][1]) 이다. 즉, 자식 노드에 대해서 멘토로 설정할지 설정하지 않을지 고려하는 것이다.

 

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

vector<vector<int>> v;
long long work[200001], n, dp[200001][2];
// dp[i][j] = i번째 노드가 j == 1이면 멘토, j == 0이면 멘토가 아닐 때 i번째 노드를 포함한 서브트리의 시너지의 최댓값

long long go(int node, int state) // 현재 노드, 멘토인지 아닌지
{
	if (dp[node][state] != -1)
	{
		return dp[node][state];
	}
	long long res = 0;
	for (auto &i : v[node])
	{
		if (state == 1) // node가 멘토일 때
		{
			// node와 i를 연결했을 때 시너지의 최댓값
			res = max(res, go(node, 0) - max(go(i, 0), go(i, 1)) + work[node] * work[i] + go(i, 0));
		}
		else // node가 멘토가 아닐 때
		{
			res += max(go(i, 1), go(i, 0));
		}
	}
	return dp[node][state] = res;
}

int main()
{
	
	cin >> n;
	v.resize(n + 1);
	memset(dp, -1, sizeof(dp));
	for (int i = 2; i <= n; i++)
	{
		int a;
		cin >> a;
		v[a].push_back(i);
	}

	for (int i = 1; i <= n; i++)
	{
		cin >> work[i];
	}

	cout << max(go(1, 1), go(1, 0));
}