「bzoj3522」/「bzoj4543」Hotel - 树形动规+长链剖分

题目大意

    求一棵树上三点距离两两相等的三元组对数。
    $n\le5000$，$n\le10^5$（加强版）

---

题目分析

$n\le5000$

设$f[i,j]$表示$i$子树中距离$i$为$j$的点的个数。
设$g[i,j]$表示$i$子树中存在两点距离它们的$lca$为$d$，而$lca$距离$i$为$d-j$。

如图，$g[i,j]$只需找到一个该子树外的$f[i,j]$即可凑成一组合法答案。

状态转移方程如下：

$$f[x,j]+=f[y,j-1]$$ $$g[x,j]+=g[y,j+1]+f[x,j]\times f[y,j]$$ $$ans+=f[x,j]\times g[x,j]$$

都很好理解。
为了防止记重，我们采用逐个合并的方式统计答案与更新状态。

$$ans+=f[x,j]\times g[y,j+1]+g[x,j]\times f[y,j-1]$$ $$g[x,j+1]+=f[x,j+1]\times f[y,j]$$ $$g[x,j-1]+=g[y,j]$$ $$f[x,j+1]+=f[y,j]$$

时间复杂度$O(n^2)$。

$n\le10^5$

发现状态只以深度为下标，故可以使用长链剖分进行优化。
用指针的方法优化重儿子转移与节省空间。

值得一提的是，本题并没有用到中间过程的状态，因为答案在中间过程已经统计，故不存在学习笔记中提到的覆盖问题。

---

代码

#include<algorithm>
#include<iostream>
#include<iomanip>
#include<cstring>
#include<cstdlib>
#include<climits>
#include<vector>
#include<cstdio>
#include<cmath>
#include<queue>
using namespace std;

typedef long long LL;

inline const int Get_Int() {
	int num=0,bj=1;
	char x=getchar();
	while(x<'0'||x>'9') {
		if(x=='-')bj=-1;
		x=getchar();
	}
	while(x>='0'&&x<='9') {
		num=num*10+x-'0';
		x=getchar();
	}
	return num*bj;
}

const int maxn=100005;

int n,Len[maxn],Son[maxn];
LL tmp[maxn*5],*f[maxn],*g[maxn],*now=tmp,ans=0;
vector<int> edges[maxn];

void AddEdge(int x,int y) {
	edges[x].push_back(y);
}

void Dfs(int Now,int fa,int depth) {
	Len[Now]=0;
	for(int Next:edges[Now]) {
		if(Next==fa)continue;
		Dfs(Next,Now,depth+1);
		if(Len[Next]>Len[Son[Now]]) {
			Son[Now]=Next;
			Len[Now]=Len[Next]+1;
		}
	}
}

void TreeDp(int Now,int fa) {
	if(Son[Now]) {
		f[Son[Now]]=f[Now]+1;
		g[Son[Now]]=g[Now]-1;
		TreeDp(Son[Now],Now);
	}
	f[Now][0]=1;
	ans+=g[Now][0];
	for(int Next:edges[Now]) {
		if(Next==fa||Next==Son[Now])continue;
		f[Next]=now,now+=(Len[Next]<<1)+2,g[Next]=now,now+=(Len[Next]<<1)+2;
		TreeDp(Next,Now);
		for(int j=0; j<=Len[Next]; j++) {
			if(j)ans+=f[Now][j-1]*g[Next][j];
			ans+=g[Now][j+1]*f[Next][j];
		}
		for(int j=0; j<=Len[Next]; j++) {
			g[Now][j+1]+=f[Now][j+1]*f[Next][j];
			if(j)g[Now][j-1]+=g[Next][j];
			f[Now][j+1]+=f[Next][j];
		}
	}
}

int main() {
	n=Get_Int();
	for(int i=1; i<n; i++) {
		int x=Get_Int(),y=Get_Int();
		AddEdge(x,y);
		AddEdge(y,x);
	}
	Len[0]=-1;
	Dfs(1,0,1);
	f[1]=now,now+=(Len[1]<<1)+2,g[1]=now,now+=(Len[1]<<1)+2; //前后分配内存，2倍!! 
	TreeDp(1,0);
	printf("%lld\n",ans);
	return 0;
}