histcat

Problem#

A certain university has $n$ staff members, numbered from $1$ to $n$.

There is a hierarchical relationship among them, meaning their relationships resemble a tree with the principal as the root, where the parent node is the direct supervisor of the child nodes.

Now there is an anniversary celebration, and inviting each staff member will increase the happiness index by $r_i$. However, if a staff member's direct supervisor attends the ball, that staff member will definitely not attend the ball.

Therefore, please write a program to calculate which staff members can be invited to maximize the happiness index, and find the maximum happiness index.

For $100%$ of the data, it is guaranteed that $1 \leq n \leq 6 \times 10^3$, $-128 \leq r_i \leq 127$, $1 \leq l, k \leq n$, and the given relationships will definitely form a tree.

Solution#

Let f[u][1/0] represent the maximum value obtained by selecting/not selecting u as part of the subtree.
Then classify accordingly.

Errors#

1. The add function nxt was written incorrectly.
2. f[u][1] should include r.

Code#

#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;
}