P1352 上司のいない舞踏会

問題#

ある大学には $n$ 人の職員がいて、番号は $1\ldots n$ です。

彼らの間には従属関係があり、つまり彼らの関係は校長を根とした木のようになっており、親ノードは子ノードの直接の上司です。

現在、周年記念パーティーがあり、パーティーに招待される職員が来るたびに一定の幸福指数 $r_i$ が増加しますが、もしある職員の直接の上司がパーティーに参加する場合、その職員はどうしてもパーティーに参加しないことになります。

したがって、どの職員を招待すれば幸福指数が最大になるかを計算するプログラムを作成し、最大の幸福指数を求めてください。

$100%$ のデータに対して、$1\leq n \leq 6 \times 10^3$、$-128 \leq r_i\leq 127$、$1 \leq l, k \leq n$ であり、与えられた関係は必ず木になります。

解答#

f[u][1/0] を u を根とする部分木で選択 / 未選択した場合の最大価値を表すとします。
それから分類すればよいです。

エラー#

add 関数の nxt が間違っている
f[u][1] に r を加える必要がある

コード#

#include<bits/stdc++.h>

using namespace std;

const int N = 6e3 + 100; 

int read()
{
	int f = 1, x = 0;
	char ch = getchar();
	while(ch < '0' || ch > '9')
	{
		if(ch == '-') f = -1;
		ch = getchar();
	}

	while(ch >= '0' && ch <= '9')
	{
		x = x * 10 + ch - '0';
		ch = getchar();
	}
	return f * x;
}

int head[N], to[N], nxt[N], cnt = 1;
int n;
int r[N], in_edge[N], root;
int f[N][2];


void add(int x, int y)
{
	to[++cnt] = y;
	nxt[cnt] = head[x];
	head[x] = cnt;
}

void dfs(int u)
{
	for(int i = head[u]; i; i = nxt[i])
	{
		int v = to[i];
		dfs(v);
		f[u][1] += f[v][0];
		f[u][0] += max(f[v][0], f[v][1]);
	}
	f[u][1] += r[u];
	if(f[u][1] == 0 && f[u][0] == 0)
	{
		f[u][1] = r[u];
	}
	return;
}

int main()
{
	n = read();

	for(int i = 1;i <= n;i++)
	{
		r[i] = read();
	}
	int x, y;
	for(int i = 1;i < n;i++)
	{
		x = read(), y = read();
		add(y, x);
		in_edge[x] ++;
	}

	for(int i = 1;i <= n;i++)
	{
		if(in_edge[i] == 0)
		{
			root = i;
			break;
		}
	}

	dfs(root);

	printf("%d", max(f[root][1], f[root][0]));
	return 0;
}