P5904 [POI2014]HOT-Hotels 题解
原题 \(n\leq 5000\)，加强版 \(n\leq 10^5\)，实际上能做 \(n\leq 10^6\)。
长链剖分。做这种很多点满足限制的题，套路差不多。设 \(f_{u,i}\) 表示 \(u\) 子树内到 \(u\) 距离为 \(i\) 的点数，\(g_{u,i}\) 表示 \(u\) 子树内的点对 \((x,y)\) 数量，满足再加一个到 \(u\) 距离为 \(i\) 的点 \(z\)，\((x,y,z)\) 可以形成题目要求的三元组。
转移的式子大概长这样：
\[f_{u,i}=\sum_{v\in son(u)}f_{v,i-1}
\]
\[g_{u,i}=\sum_{v\in son(u)} f_{v,i-1}\cdot f_{u,i}+\sum_{v\in son(u)} g_{v,i+1}
\]
\[ans=\sum_u \big(\sum_{v\in son(u)}\sum_i f_{u,i-1}\cdot g_{v,i}+\sum_{v\in son(u)}\sum_i g_{u,i+1}\cdot f_{v,i}\big)
\]
使用长剖优化，动态开内存给长链顶部，每个点直接继承重儿子信息，把轻儿子合并，复杂度 \(O(n)\)。
注意动态开内存的大小，如果有 \(-1\) 的数组，需要在前后都留下足够的空间。
点击查看代码
const int N=1e5+13;
int n,d[N],son[N];
ll buf1[N<<3],buf2[N<<4];
ll *f[N],*g[N],*nowf=buf1,*nowg=buf2;
std::vector<int> e[N];
ll ans;
void dfs(int u,int fa){
for(auto v:e[u]){
if(v==fa) continue;
dfs(v,u);
if(d[v]>d[son[u]]) son[u]=v;
}
d[u]=d[son[u]]+1;
}
void dp(int u,int fa){
if(son[u]) f[son[u]]=f[u]+1,g[son[u]]=g[u]-1,dp(son[u],u);
f[u][0]=1,ans+=g[u][0];
for(auto v:e[u]){
if(v==fa||v==son[u]) continue;
f[v]=nowf,nowf+=d[v]+4;nowg+=(d[v]<<1)+10,g[v]=nowg;
dp(v,u);
for(int i=0;i<d[v];++i) ans+=(i?f[u][i-1]*g[v][i]:0)+g[u][i+1]*f[v][i];
for(int i=0;i<=d[v];++i) g[u][i]+=(i?f[v][i-1]*f[u][i]:0)+(i<d[v]?g[v][i+1]:0);
for(int i=1;i<=d[v];++i) f[u][i]+=f[v][i-1];
}
}
int main(){
// freopen("4.in","r",stdin);
// freopen("P5904.out","w",stdout);
read(n);
for(int i=1;i<n;++i){
int u,v;read(u),read(v);
e[u].pb(v),e[v].pb(u);
}
dfs(1,0);
f[1]=nowf,nowf+=d[1]+4;
nowg+=(d[1]<<1)+4,g[1]=nowg;
dp(1,0);
println(ans);
return 0;
}
