分析
题目中说只有距离为2的点才能产生联合权值,让我们画画图理解一下
这个时候我们发现所有的联合权值之和其实就是以每个点为根,其所有子树的权值两两的乘积之和。假设其中一个点为j,他对答案的贡献就是当前根的出j以外的儿子的权值总和乘节点j的权值
for (int i=head[x];i;i=e[i].nex)
{
int v=e[i].to;
ans=(ans+a[v]*(sum-a[v]+MOD)%MOD)%MOD;
//a[v]是每个节点的权值,sum是当前根所有子节点的权值和
}
至于最大值,还记得怎么求树的直径的吗,同理,我们记录一个ma与m,表示次大值与最大值,用一个变量记录ma*m的最大值就好了
ll sum=0,ma=0,m=0;
for (int i=head[x];i;i=e[i].nex)
{
int v=e[i].to;
if(a[v]>ma) m=ma,ma=a[v];
else if(a[v]>m) m=a[v];
sum=(sum+a[v])%MOD;
}
maxn=max(maxn,ma*m);
代码
#include<bits/stdc++.h>
#define ll long long
using namespace std;
const int N=2e5+5;
const int MOD=1e4+7;
ll n,cnt,maxn,ans;
int a[N],head[N],f[N];
bool flag[N];
struct nk
{
int to,nex;
}e[N<<1];
inline void add(int x,int y)
{
e[++cnt].nex=head[x];e[cnt].to=y;head[x]=cnt;
}
inline void dfs(int x)
{
ll sum=0,ma=0,m=0;
for (int i=head[x];i;i=e[i].nex)
{
int v=e[i].to;
if(a[v]>ma) m=ma,ma=a[v];
else if(a[v]>m) m=a[v];
sum=(sum+a[v])%MOD;
}
maxn=max(maxn,ma*m);
for (int i=head[x];i;i=e[i].nex)
{
int v=e[i].to;
ans=(ans+a[v]*(sum-a[v]+MOD)%MOD)%MOD;
}
}
int main()
{
scanf("%lld",&n);
for (int i=1;i<n;i++)
{
int u,v;
scanf("%d%d",&u,&v);
add(u,v),add(v,u);
}
for (int i=1;i<=n;i++) scanf("%d",&a[i]);
for (int i=1;i<=n;i++) dfs(i);
printf("%lld %lld\n",maxn,ans%MOD);
return 0;
}