树上dp的好题，Piggy and Trees思路做法详细解析

题目链接:

• link，点击这里喵。

题意简述:

• 给出一棵树，取一点 $i$，并另取两点 $u$ 和 $v$，在 $u$ 和 $v$ 简单路径上任取一点使这点与 $i$ 的距离最小，计算出这样所有的组合的距离的总和。

思路:

• 看眼数据范围，总点数 $n\le 2\times 10^5$，考虑线性算法，有关于树，首先考虑树上 dp。
• 这道题题面很具有迷惑性，她把 $\sum _{u=1}^n\sum _{v=u+1}^n\sum _{i=1}^{n}f(u,v,i)$ 中的 $\sum\limits_{i = 1}^n $ 放在了后面，这就导致如果思路想到枚举 $u$ 和 $v$，就很难 dp 了，根据 $\sum$ 的性质，我们把 $\sum\limits_{i = 1}^n $ 提到外面去，原式就变成了 $\sum _{i=1}^n\sum _{u=1}^n\sum _{v=u+1}^{n}f(u,v,i)$，思路就简洁明了了，记录每个点的答案，最后相加即可。
• 那么 $O(n)$ 做法考虑换根 dp，这里假设以点 $1$ 为出发点。

做法:

• 定义 $l{j}_i$ 为点 $i$ 及其子树所包含的点对数量，$si{z}_i$ 为点 $i$ 及其子树的大小，$si{z}_i$ 为以点 $i$ 为 $\sum _{u=1}^n\sum _{v=u+1}^{n}f(u,v,i)$ 的答案。

  dfs1:

  • $si{z}_{now}=\sum _{to}si{z}_{to}$ （$to$ 为 $now$ 的子节点）//这个应该不必多言
  • $l{j}_{now}=\sum _{to}l{j}_{to}+\sum _{to}\sum _{to2\ne to}si{z}_{to}\times si{z}_{to2}+(si{z}_{now}-1)$（$to，to2$ 为 $now$ 的子节点）
    $\sum _{to}l{j}_{to}$ 应该很容易理解，加上子树点对数。
    $\sum _{to}\sum _{to2\ne to}si{z}_{to}\times si{z}_{to2}$ 就是乘法原理，每个子树有多少个可以选择的，互相乘起来就好，这里这么写方便，实际应当去除重复计算的。
    $(si{z}_{now}-1)$ 为当前节点与子树上每个节点所构成的点对。
  • $an{s}_{now}=\sum _{to}(an{s}_{to}+l{j}_{to})$
    $l{j}_{to}$ 便是原先所有点对数量，由于回到了 $now$，单子树整个的花费比原来节点多了 $l{j}_{to}\times 1$，总花费便是 $(an{s}_{to}+l{j}_{to})$
  • 代码实现：

  void dfs1(int now,int fa){
     siz[now]=1;
     for(auto to: G[now]){
         if(to==fa){continue;}
         dfs1(to,now);
         siz[now]+=siz[to];
         ans[now]+=ans[to]+lj[to];
         lj[now]+=lj[to];
         mod(lj[now]);mod(ans[now]);
     }
     lnt sum=0;
     for(auto to: G[now]){if(to==fa){continue;}sum+=siz[to];}//这里进行了优化，把O(n^2)的枚举优化成O(n)
     for(int i=0;i<G[now].size();++i){                       //其实就是计算了个总和而已啦
         int to=G[now][i];
         if(to==fa){continue;}
         sum-=siz[to];
         lj[now]+=siz[to]*sum;
     }
     lj[now]+=siz[now]-1;
     lj[now]%=MOD;
  }
  

  dfs2:

  那么有了公式就好推了，仍然从 $1$ 出发。
  枚举 $now$ 的子节点 $to$，由先前所推的公式易得：

  • 不算子树 $to$ 时，$l{j}_{nownew}=l{j}_{now}-l{j}_{to}-si{z}_{to}\times (si{z}_{now}-si{z}_{to}-1)-si{z}_{to}$
  • 同理可知不算子树 $to$ 时，$an{s}_{nowtemp}=an{s}_{now}-an{s}_{to}-l{j}_{to}$
  • 故可知，加入 $now$ 为子树的 $to$ 有：
    $l{j}_{tonew}=(si{z}_{now}-si{z}_{to})\times (si{z}_{to}-1)+l{j}_{nownew}+si{z}_{now}-si{z}_{to}$
    $an{s}_{tonew}=an{s}_{to}+an{s}_{now}-an{s}_{to}-l{j}_{to}+l{j}_{nownew}$
    $si{z}_{to}=n$
  • 代码实现：

  void dfs2(int now,int fa){
     anst+=ans[now];
     mod(anst);
     for(auto to: G[now]){
         if(to==fa){continue;}
         lnt A=lj[now]-lj[to]-siz[to]*(siz[now]-siz[to]-1)-siz[to];
         mod(A);
         ans[to]+=ans[now]-ans[to]-lj[to]+A;
         mod(ans[to]);
         lj[to]+=(siz[now]-siz[to])*(siz[to]-1)+A+siz[now]-siz[to];
         mod(lj[to]);
         siz[to]=n;
         dfs2(to,now);
     }
  }
  

小贴士:

• 记得开 longlong，乘法会溢出。
• 记得判断父亲节点，不然会 RE。

时间复杂度分析

• 每个点只在 dfs1 和 dfs2 中分别出现一遍，故时间复杂度时线性 $O(n)$ 的。

代码:

#include <stdio.h>
#include <ctype.h>
#include <algorithm>
#include <string.h>
#include <vector>
#define lnt long long
#define inf 0x3f3f3f3f
using namespace std;
int xx;char chh,ff;int read(){
	xx=ff=0;while(!isdigit(chh)){if(chh=='-'){ff=1;}chh=getchar();}
	while(isdigit(chh)){xx=(xx<<1)+(xx<<3)+chh-'0';chh=getchar();}return ff? -xx: xx;
}
const int N=1e6,MOD=1e9+7;
lnt n,ans[N],anst,siz[N],lj[N];//
vector<lnt> G[N];
void dfs1(int,int);
void dfs2(int,int);
void mod(lnt &x){x%=MOD;}
int main(){
	n=read();
	for(int i=1;i<n;++i){
		int x=read(),y=read();
		G[x].push_back(y);
		G[y].push_back(x);
	}
	dfs1(1,0);
	dfs2(1,0);
	printf("%lld",(anst+MOD)%MOD);
	
	return 0;
}
void dfs1(int now,int fa){
	siz[now]=1;
	for(auto to: G[now]){
		if(to==fa){continue;}
		dfs1(to,now);
		siz[now]+=siz[to];
		ans[now]+=ans[to]+lj[to];
		lj[now]+=lj[to];
		mod(lj[now]);mod(ans[now]);
	}
	lnt sum=0;
	for(auto to: G[now]){if(to==fa){continue;}sum+=siz[to];}
	for(int i=0;i<G[now].size();++i){
		int to=G[now][i];
		if(to==fa){continue;}
		sum-=siz[to];
		lj[now]+=siz[to]*sum;
	}
	lj[now]+=siz[now]-1;
	lj[now]%=MOD;
}
void dfs2(int now,int fa){
	anst+=ans[now];
	mod(anst);
	for(auto to: G[now]){
		if(to==fa){continue;}
		lnt A=lj[now]-lj[to]-siz[to]*(siz[now]-siz[to]-1)-siz[to];
		mod(A);
		ans[to]+=ans[now]-ans[to]-lj[to]+A;
		mod(ans[to]);
		lj[to]+=(siz[now]-siz[to])*(siz[to]-1)+A+siz[now]-siz[to];
		mod(lj[to]);
		siz[to]=n;
		dfs2(to,now);
	}
}