[bzoj4543][bzoj3522][DP][长链剖分]Hotel&Hotel加强版

Description

有一个树形结构的宾馆，n个房间，n-1条无向边，每条边的长度相同，任意两个房间可以相互到达。吉丽要给他的三个妹子各开（一个）房（间）。三个妹子住的房间要互不相同（否则要打起来了），为了让吉丽满意，你需要让三个房间两两距离相同。
有多少种方案能让吉丽满意？

Input

第一行一个数n。 接下来n-1行，每行两个数x,y，表示x和y之间有一条边相连。

Output

让吉丽满意的方案数。

Sample Input

7

1 2

5 7

2 5

2 3

5 6

4 5

Sample Output

5

HINT

【样例解释】

{1,3,5},{2,4,6},{2,4,7},{2,6,7},{4,6,7}

【数据范围】

n≤5000

题解

首先有个性质
这三个点的路径会交于一点，三个点到这个点的距离都相等且位于同一棵子树
那么$N^2$的就可以直接做了…
这个$O(N)$的有点神仙…
设一个dp状态
$f[i][j]$表示$i$点子树中有多少距离为$j$的点
$g[i][j]$表示$i$点子树中，有多少对，在$i$的子树外面找一个距离为$j$的点就能构成一个合法三元组的数量
每扫到一棵子树$y$
可以这样累加答案
$$ans+=\sum f[x][j]\ast g[y][j+1]+g[x][j]\ast f[y][j-1]$$
然后更新这样更新
$$g[x][i]+=f[x][i]\ast f[y][i-1]$$
$$f[x][i]+=f[y][i-1]$$
$$g[x][i]+=g[y][i+1]$$
注意到这样是$\sum maxdep[i]$的
考虑如何优化
我们发现，每个点的第一次转移其实就是找一个儿子，让他的$g$数组左移一位，$f$数组右移一位
指针移动即可
这样可以节省$\sum maxdep[son[i]]$的复杂度
长链剖分，每个点选择它的重儿子做更新
其他儿子暴力更新
单考虑重链上的转移，一次是$O(1)$的，所以总复杂度不会超过$O(N)$
再考虑轻边，每条重链只会在头的位置被计算完整个重链的长度
显然重链没有重叠
所以轻边的转移总复杂度也是$O(N)$的
于是就是$O(N)$的优秀复杂度了…
空间的话…学习了一下别人的开法毕竟不会指针数组这个东西
upd:
考试忘记怎么开了
那就手写了一个
感觉还挺好写的qaq

#include<cstdio>
#include<cstring>
#include<cstdlib>
#include<algorithm>
#include<cmath>
#include<queue>
#include<vector>
#include<ctime>
#include<map>
#include<bitset>
#include<set>
#define LL long long
#define mp(x,y) make_pair(x,y)
#define pll pair<long long,long long>
#define pii pair<int,int>
#define f(x,i) f1[beginf[x]+i]
#define g(x,i) f2[beging[x]+i]
using namespace std;
inline int read()
{
	int f=1,x=0;char ch=getchar();
	while(ch<'0'||ch>'9'){if(ch=='-')f=-1;ch=getchar();}
	while(ch>='0'&&ch<='9'){x=x*10+ch-'0';ch=getchar();}
	return x*f;
}
int stack[20];
inline void write(LL x)
{
	if(x<0){putchar('-');x=-x;}
    if(!x){putchar('0');return;}
    int top=0;
    while(x)stack[++top]=x%10,x/=10;
    while(top)putchar(stack[top--]+'0');
}
inline void pr1(int x){write(x);putchar(' ');}
inline void pr2(LL x){write(x);putchar('\n');}
const int MAXN=100005;
struct edge{int x,y,next;}a[2*MAXN];int len,last[MAXN];
void ins(int x,int y){len++;a[len].x=x;a[len].y=y;a[len].next=last[x];last[x]=len;}

LL f1[MAXN*2],f2[MAXN*2];
int beginf[MAXN],beging[MAXN],T;

int dep[MAXN],son[MAXN],maxdep[MAXN],fa[MAXN];
void pre_tree_node(int x)
{
	maxdep[x]=dep[x];
	for(int k=last[x];k;k=a[k].next)
	{
		int y=a[k].y;
		if(y!=fa[x])
		{
			fa[y]=x;dep[y]=dep[x]+1;
			pre_tree_node(y);
			if(maxdep[y]>maxdep[son[x]])son[x]=y;
			maxdep[x]=max(maxdep[x],maxdep[y]);
		}
	}
}
void pre_tree_edge(int x,int u1,int u2)
{
	beginf[x]=u1;beging[x]=u2;
	if(son[x])pre_tree_edge(son[x],u1+1,u2-1);
	for(int k=last[x];k;k=a[k].next)
	{
		int y=a[k].y;
		if(y!=fa[x]&&y!=son[x])
		{
			int lin=T;
			T+=(maxdep[y]-dep[y]+1)*2;
			pre_tree_edge(y,lin+1,T-(maxdep[y]-dep[y]+1)+1);
		}
	}
}
int n;
LL ans;
void dp(int x,int fa)
{
	f(x,0)=1;
	if(son[x])dp(son[x],x);
	ans+=g(x,0);
	for(int k=last[x];k;k=a[k].next)
	{
		int y=a[k].y;
		if(y!=fa&&y!=son[x])
		{
			dp(y,x);
			int ln=maxdep[y]-dep[y];
			for(int i=0;i<=ln;i++)ans+=f(y,i)*g(x,i+1);
			for(int i=1;i<=ln;i++)ans+=g(y,i)*f(x,i-1);
			
			for(int i=0;i<=ln;i++)g(x,i+1)+=f(y,i)*f(x,i+1);
			for(int i=1;i<=ln;i++)g(x,i-1)+=g(y,i);
			
			for(int i=0;i<=ln;i++)f(x,i+1)+=f(y,i);
		}
	}
}
int main()
{
	n=read();
		for(int i=1;i<=2*n;i++)f1[i]=f2[i]=0;
		len=0;memset(last,0,sizeof(last));
		for(int i=1;i<n;i++)
		{
			int x=read(),y=read();
			ins(x,y);ins(y,x);
		}
		memset(fa,0,sizeof(fa));memset(son,0,sizeof(son));
		pre_tree_node(1);
		T=2*(maxdep[1]-dep[1]+1);
		pre_tree_edge(1,1,T-(maxdep[1]-dep[1]+1)+1);
		ans=0;
		dp(1,0);
		pr2(ans);
//	}
	return 0;
}

#include<cstdio>
#include<cstring>
#include<cstdlib>
#include<algorithm>
#include<cmath>
#include<queue>
#include<vector>
#include<ctime>
#include<map>
#include<bitset>
#define LL long long
#define mp(x,y) make_pair(x,y)
using namespace std;
inline int read()
{
	int f=1,x=0;char ch=getchar();
	while(ch<'0'||ch>'9'){if(ch=='-')f=-1;ch=getchar();}
	while(ch>='0'&&ch<='9'){x=x*10+ch-'0';ch=getchar();}
	return x*f;
}
int stack[20];
inline void write(LL x)
{
    if(!x){putchar('0');return;}
    int top=0;
    while(x)stack[++top]=x%10,x/=10;
    while(top)putchar(stack[top--]+'0');
}
inline void pr1(int x){write(x);putchar(' ');}
inline void pr2(LL x){write(x);putchar('\n');}
struct node{int y,next;}a[210000];int len,last[110000];
void ins(int x,int y){len++;a[len].y=y;a[len].next=last[x];last[x]=len;}
LL spa[100005*10];
LL *f[100005],*g[100005],*now=spa+100005,ans;
int dep[100005],son[100005];
void init(int x,int fa)
{
	for(int k=last[x];k;k=a[k].next)
	{
		int y=a[k].y;
		if(y!=fa)
		{
			init(y,x);
			if(dep[y]>dep[son[x]])son[x]=y;
		}
	}
	dep[x]=dep[son[x]]+1;
}
void newnode(int x)
{
	f[x]=now;now=now+2*dep[x]+1;
	g[x]=now;now=now+2*dep[x]+1;
}
void dp(int x,int fa)
{
	f[x][0]=1;
	if(son[x])
	{
		f[son[x]]=f[x]+1;g[son[x]]=g[x]-1;
		dp(son[x],x);
		ans+=g[son[x]][1];
	}
	for(int k=last[x];k;k=a[k].next)
	{
		int y=a[k].y;
		if(y!=fa&&y!=son[x])
		{
			newnode(y);dp(y,x);
			for(int j=dep[y];j>=0;j--)
			{
				if(j)ans=ans+g[y][j]*f[x][j-1];
				ans=ans+g[x][j+1]*f[y][j];
				g[x][j+1]+=f[x][j+1]*f[y][j];
			}
			for(int j=dep[y];j>=0;j--)
			{
				if(j)g[x][j-1]+=g[y][j];
				f[x][j+1]+=f[y][j];
			}
		}
	}
}
int n;
int main()
{
	n=read();
	for(int i=1;i<n;i++)
	{
		int x=read(),y=read();
		ins(x,y);ins(y,x);
	}
	init(1,0);
	newnode(1);
	dp(1,0);
	pr2(ans);
	return 0;
}