[BZOJ4543/3522][POI2014]Hotel加强版(长链剖分+树上DP)

题意简述

给出一棵 $n$ 个点的无根树，请在这棵树上选三个互不相同的节点，使得这个三个节点两两之间距离相等，输出方案数即可。
$subtask1:\sum n\le 500$
$subtask2:\sum n\le 5000$
$subtask3:\sum n\le 500000$

题解

subtask 1:明显是送分点，直接用Floyd跑都可以。
subtask 2:考虑如何写一个树上DP来解决这个问题。
$f[u][i]$表示$u$子树里距离$u$为$i$的节点个数
$g[u][i]$表示$u$子树里两个点$x,y$到其$LCA$的距离都是$d$,$LCA$到$u$的距离是$d?i$的方案数
有
$g[u][i]={\sum }_{v}g[v][j+1]+f[u][i]?f[v][i?1]$
$f[u][i]={\sum }_{v}f[v][i?1]$
注意到当一个节点只有一个时:
$g[u][i]=g[v][j+1]$
$f[u][i]=f[v][i?1]$
可以用指针转移实现$O(1)$。
每个节点选择深度最大的那一个点转移即可。
注意到每一个点都属于一个长链，而一条长链中的所有节点都可以$O(1)$转移，只有在链顶才会$O(该链长度)$地暴力转移。
因此复杂度:$O(2\sum 链长度)=O(n)$

实现

#include<iostream>
#include<cstdio>
#include<algorithm>
#include<cmath>
#include<queue>
#include<cstring>
#include<ctime>
using namespace std;
#define MAXM 50000
#define MAXK 1000
#define LL long long
int n,x,y,ecnt;
struct edge
{
    int nxt,to;
}e[MAXM*2+5];
LL space[MAXM*10+1];
LL *f[MAXM+5],*dp[MAXM+5],*tot=space+MAXM;
int dep[MAXM+5],id[MAXM+5],pnt;
int head[MAXM+5],fa[MAXM+5];
LL ans=0;
inline int read()  
{
      int x=0,f=1;char s=getchar();  while(s<'0'||s>'9'){
    if(s=='-')f=-1;s=getchar();}  while(s>='0'&&s<='9'){
    x=x*10+s-'0';s=getchar();}  return x*f;  
}
void add(int u,int v)
{
    e[++ecnt]=(edge){
    head[u],v};head[u]=ecnt;e[++ecnt]=(edge){
    head[v],u};head[v]=ecnt;
}
void new_node(int x)
{
    dp[x]=tot;tot=tot+dep[x]*2+1;pnt+=dep[x]*4+2;f[x]=tot;tot=tot+dep[x]*2+1;
}
void dfs1(int x)
{
    dep[x]=1;id[x]=0;for(int i=head[x];i!=-1;i=e[i].nxt){
    int xnt=e[i].to;if(xnt==fa[x])continue;fa[xnt]=x;dfs1(xnt);if(dep[x]<dep[xnt]+1)dep[x]=dep[xnt]+1,id[x]=xnt;}
}
void dfs2(int x)
{
    dp[x][0]=1;if(id[x]){
    dp[id[x]]=dp[x]+1;f[id[x]]=f[x]-1;//神奇指针dfs2(id[x]);ans+=f[id[x]][1];}for(int i=head[x];i!=-1;i=e[i].nxt){
    int xnt=e[i].to;if(xnt==fa[x]||xnt==id[x])continue;new_node(xnt);dfs2(xnt);for(int j=dep[xnt];j>=0;j--){
    ans+=dp[x][j]*f[xnt][j+1];if(j>0){
    ans+=f[x][j]*dp[xnt][j-1];f[x][j]+=dp[x][j]*dp[xnt][j-1];dp[x][j]+=dp[xnt][j-1];}f[x][j]+=f[xnt][j+1];}}
}
int main()
{
    while(~scanf("%d",&n)&&n){
    pnt=MAXM;tot=space+MAXM;ecnt=0;ans=0;for(int i=1;i<=n;i++)head[i]=-1;for(int i=1;i<n;i++){
    x=read(),y=read();add(x,y);}dfs1(1);new_node(1);dfs2(1);printf("%lld\n",ans);for(int i=MAXM;i<=pnt;i++)space[i]=0;}
}