633F - The Chocolate Spree - CodeForces Solution

#include<bits/stdc++.h>
#define LL long long
#define N 100010
using namespace std;
int n;
int w[N];
LL ans;
LL f[N][2];//最大的链(0)，整段(1) 
vector<int>ve[N];
void dfs1(int u,int fa)
{
	LL k1=0,k2=0;
	for(auto &v:ve[u])
	{
		if(v==fa)	continue;
		dfs1(v,u);
		f[u][0]=max(f[u][0],f[v][0]);
		if(f[v][1]>=k1)
		{
			k2=k1;
			k1=f[v][1];
		}
		else if(f[v][1]>=k2)	k2=f[v][1];
	}
	f[u][0]=max(f[u][0],k1+k2+w[u]);
	f[u][1]=w[u]+k1;
}
void dfs2(int u,int fa,LL t1,LL t2)//t1整段 t2链
{
	LL k1=0,k2=0,p1=0,p2=0,p3=0;//k整段 p链
	for(auto &v:ve[u])
	{
		if(v==fa)	continue;
		if(f[v][0]>=k1)
		{
			k2=k1;
			k1=f[v][0];
		}
		else if(f[v][0]>=k2)	k2=f[v][0];
		if(f[v][1]>=p1)
		{
			p3=p2;
			p2=p1;
			p1=f[v][1];
		}
		else if(f[v][1]>=p2)
		{
			p3=p2;
			p2=f[v][1];
		}
		else if(f[v][1]>=p3)	p3=f[v][1];
	}
	for(auto &v:ve[u])
	{
		if(v==fa)	continue;
		LL newt1=max(f[v][0]==k1?k2:k1,t1);
		LL newt2=max(f[v][1]==p1?p2:p1,t2)+w[u];
		if(f[v][1]==p1)	newt1=max(newt1,t2+p2+w[u]);
		else newt1=max(newt1,t2+p1+w[u]);
		if(f[v][1]==p1)	newt1=max(newt1,p2+p3+w[u]);
		else if(f[v][1]==p2)	newt1=max(newt1,p1+p3+w[u]);
		else	newt1=max(newt1,p1+p2+w[u]);
		ans=max(ans,newt1+f[v][0]);
		dfs2(v,u,newt1,newt2);
	}
}
int main()
{
	cin>>n;
	for(int i=1;i<=n;i++)	cin>>w[i];
	for(int i=1;i<n;i++)
	{
		int u,v;
		cin>>u>>v;
		ve[u].push_back(v);
		ve[v].push_back(u);
	}
	dfs1(1,0);
	dfs2(1,0,0,0);
	cout<<ans;
	return 0;
}