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

题目来源：洛谷P1351

思路

由题意可得图为一棵树

在一棵树上距离为2的两个点有两种情况

• 当前点与其爷爷
• 当前点的两个儿子

当情况为当前点与其爷爷时比较好操作

只需要在传递时不仅传递父亲 还传递爷爷即可

当情况为两个儿子时

其实我们只需要的是所有儿子中比较大的两个

所以我们在遍历当前点的所有儿子时取出最大值和次大值即可

最后的最大值为两种情况中的最大值

最后的答案为当前点的所有儿子权值的平方减去每个儿子本身的平方（容斥原理）+2*最大*次大（乘以2是因为每种情况有正反两种）

代码

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