树形DP CH5401没有上司的舞会

树形dp入门题

 

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
//typedef __int128 LL;
//typedef unsigned long long ull;
#define F first
#define S second
typedef long double ld;
typedef pair<int,int> pii;
typedef pair<ll,ll> pll;
typedef pair<ld,ld> pdd;
const ld PI=acos(-1);
const ld eps=1e-9;
//unordered_map<int,int>mp;
#define ls (o<<1)
#define rs (o<<1|1)
#define pb push_back
const int seed=131;
const int M = 6000+7;
struct node
{
	int nxt,to,val;
}ee[M*2];
int head[M*2];
int cnt;
void add(int x,int y)
{
	ee[++cnt].to=y;
	ee[cnt].nxt=head[x];
//	ee[cnt].val=z;
	head[x]=cnt;
}
int z[M],ru[M],chu[M];
int rt;
int dp[M][M];//以i为根节点的子树，选择i节点，快乐指数和最大是多少 
void dfs(int u,int fa)
{
	dp[u][1]=z[u];
//	dp[u][0]=0;
	for(int i=head[u];i;i=ee[i].nxt)
	{
		int v=ee[i].to;
		if(v==fa)continue;
		dfs(v,u);
		dp[u][1]+=dp[v][0];
		dp[u][0]+=max(dp[v][1],dp[v][0]);
	}
//	cout<<u<<" "<<fa<<" "<<" "<<dp[u][0]<<" "<<dp[u][1]<<endl; 
	
//	dp[u][0]=max(h0,h1);
}
int main()
{
	ios::sync_with_stdio(false);
  	cin.tie(0);
  	int n,u,v;
	cin>>n;
	for(int i=1;i<=n;i++)
	cin>>z[i];
	for(int i=1;i<n;i++)
	cin>>u>>v,add(u,v),add(v,u),chu[u]++,ru[v]++;
	cin>>u>>v;
	for(int i=1;i<=n;i++)
	{
		if(chu[i]==0)rt=i;
		if(ru[i]==0)dp[i][1]=z[i];
	}
	dfs(rt,0);
	cout<<max(dp[rt][0],dp[rt][1])<<endl;
	return 0;
}