题目链接
https://www.luogu.org/problem/P1351
题解
如果一个点对内两点的距离恰好为 \(2\)，意味着必然存在一个点同时与这两个点相邻。
于是我们可以枚举这个中转点。
对于一个中转点，其联合权值最大值必然是与其相邻的点中最大的两个权值的积。因此我们可以在 \(O(n)\) 的时间内求出最大的联合权值（树上每条边只会遍历两次）。
问题在于求出所有联合权值之和。最坏情况下（菊花图）能产生联合权值的点对的数量是 \(O(n^2)\) 的，因此直接找所有满足条件的点对是不可行的。
让我们推推式子。设与点 \(u\) 相邻的点分别为 \(v_1 ,v_2 ,\ldots ,v_n\)。显然以点 \(u\) 为中转点时，所有满足条件的联合权值之和为：
$$\sum_{i=1}^{n} \sum_{j=1}^{n} w_{v_i} \times w_{v_j} \times [i \neq j]$$
根据乘法分配律，我们把 \(w_{v_i}\) 提出来，整理一下：
$$\sum_{i=1}^{n} w_{v_i} \times (\sum_{j=1}^{n} w_{v_j} \times [i \neq j])$$
设与 \(u\) 点相邻的所有点的权值和为 \(S_u\)，带入上面的式子中，可以得到：
$$\sum_{i=1}^{n} w_{v_i} \times (S_u-w_{v_i})$$
到这里，我们已经有了一个 \(O(n)\) 算法：我们还是枚举中转点 \(u\)，遍历所有 \(v_i\) 的时候把 \(S_u\) 算出来，带入上式求和即可。
不过，我们还可以走的更远。
把上面的式子展开：
$$\sum_{i=1}^n w_{v_i} \times S_u – \sum_{i=1}^nw_{v_i}^2$$
整理后可得：
$$S_u^2 – \sum_{i=1}^nw_{v_i}^2$$
这样，我们就得到了一个非常简洁的联合权值的和的表达式。
#include <cstdio>
#include <algorithm>
#define MOD 10007
using namespace std;
struct edge
{
int v,next;
}e[400005];
int w[200005],head[200005],cnt;
void addedge(int u,int v)
{
e[++cnt].v=v;
e[cnt].next=head[u];
head[u]=cnt;
}
int main()
{
int n;
scanf("%d",&n);
for(int i=1;i<n;i++)
{
int u,v;
scanf("%d%d",&u,&v);
addedge(u,v);
addedge(v,u);
}
for(int i=1;i<=n;i++)
scanf("%d",&w[i]);
int ans=0,maxx=0;
for(int u=1;u<=n;u++)
{
int max1=0,max2=0,sum1=0,sum2=0;
for(int i=head[u];i;i=e[i].next)
{
if(w[e[i].v]>max1)max2=max1,max1=w[e[i].v];
else if(w[e[i].v]>max2)max2=w[e[i].v];
sum1=(sum1+w[e[i].v])%MOD;
sum2=(sum2+w[e[i].v]*w[e[i].v])%MOD;
}
ans=((ans+sum1*sum1-sum2)%MOD+MOD)%MOD;
maxx=max(maxx,max1*max2);
}
printf("%d %d\n",maxx,ans);
return 0;
}