【题解】洛谷P1351 [NOIP2014TG] 联合权值（树形结构+DFS）

#include<iostream>
using namespace std;
#define maxn 200020
#define mod 10007
int n,cnt,ans,Max;
int h[maxn],w[maxn];
struct Edge
{
    int to;
    int next;
}e[maxn*2];
void add(int u,int v)
{
    e[++cnt].to=v;
    e[cnt].next=h[u];
    h[u]=cnt;
}
void dfs(int u,int fa,int gr)//u为当前点 fa为父亲 gr为爷爷 
{
    int m1=0,m2=0,sum=0,num=0;//m1为最大 m2为次大 sum为所有儿子的权值 num为所有儿子本身的平方 
    for(int i=h[u];i;i=e[i].next)
    {
        int v=e[i].to;
        if(v==fa) continue;//如果是父亲就跳过 
        sum=(sum+w[v])%mod;//计算所有儿子权值 
        num=(num+w[v]*w[v]%mod)%mod;//计算儿子本身平方 
        if(w[v]>m1)//判断最大和次大 
        {    
            m2=m1;
            m1=w[v];
        }
        else if(w[v]>m2)
        m2=w[v];
        dfs(v,u,fa);
    }
    Max=max(Max,max(m1*m2,w[u]*w[gr]));//注意最大值不用取mod 
    ans=(ans+(sum*sum%mod-num+mod)%mod+w[u]*w[gr]*2%mod)%mod;//总值计算 
}
int main()
{
    cin>>n;
    for(int i=1;i<n;i++)
    {
        int x,y;
        cin>>x>>y;
        add(x,y);
        add(y,x);
    }
    for(int i=1;i<=n;i++) cin>>w[i];
    dfs(1,0,0);
    cout<<Max<<" "<<ans;
}