【题解】 P5904 [POI2014]HOT-Hotels 加强版

传送门

题意

给定一棵树，求树上存在多少个三元组 $(a,b,c)$，满足 $\mathrm{dis}(a,b)=\mathrm{dis}(a,c)=\mathrm{dis}(b,c)$。

分析

发现答案无非就是下面两种。

对于左边的情况，我们显然可以用 $f_{i,j}$ 表示 $i$ 子树内到 $i$ 距离为 $j$ 的点，但显然会算重复，例如把儿子节点的答案也算上去。

那么就先考虑右边的情况，考虑用 $g_{i,j}$ 表示 $i$ 的子树内，有多少二元组 $(a,b)$ 满足 $dis(a,lca(a,b))=dis(b,lca(a,b))=dis(i,lca(a,b))+j$:

那么右边的情况就可以看成：

这样的话，只需保证左边的节点和右边两个节点不在某一个点的同一个儿子的子树内，那么就不会被重复统计。

还有一种情况有一个点是 $now$ 本身，那么显然是 $g_{now,0}$

那么，节点 $now$ 对答案的贡献就是:
$$g_{now,0}+\sum _j\sum _{x,y\in son(now)}{f}_{x,j-1}\cdot g_{y,j+1}$$

接下来考虑如何转移 $f$ 和 $g$ 数组，

$f$ 数组转移较易，

$$f_{now,j}=\sum _{x\in son(now)}{f}_{x,j-1}$$

考虑 $g$ 数组的转移，有以下两种情况:

对于左边的情况，发现 $i$ 到 $lca(a,b)$ 的距离增加了 1, 则相应的 $j$ 减 $1$ 即可。

而右边的情况，发现 $lca(a,b)=now$ ，那么 $dis(a,lca(a,b))=dis(b,lca(a,b))=dis(now,lca(a,b))+j=j$，于是 $g$ 数组的转移如下：

$$g_{now,j}=\sum _{x,y\in son(x),x\ne y}{f}_{x,j-1}\cdot f_{y,j-1}+\sum _{x\in son(now)}{g}_{x,j+1}$$

这样我们就得到了一个在使用前缀和的前提下 $O(n^2)$ 的动态规划，但显然还可以用 长链剖分优化成 $O(n)$。

#include<bits/stdc++.h>
#define int long long
using namespace std;
const int N = 1e5 + 67;
int read(){
	int x = 0, f = 1; char ch = getchar();
	while(ch < '0' || ch > '9'){if(ch == '-') f = -f; ch = getchar();}
	while(ch >= '0' && ch <= '9'){x = (x << 1) + (x << 3) + (ch ^ 48); ch = getchar();}
	return x * f;
}
int n, ans;
int tot, Head[N], to[N << 1], Next[N << 1];
int mxd[N], son[N], *f[N], *g[N], suf[N << 2], *cur = suf;
void add(int u, int v){
	to[++tot] = v, Next[tot] = Head[u], Head[u] = tot;
}
void dfs1(int x, int fa){
	for(int i = Head[x]; i; i = Next[i]){
		int y = to[i]; if(y == fa) continue;
		dfs1(y, x); if(mxd[y] > mxd[son[x]]) son[x] = y;
	}
	mxd[x] = mxd[son[x]] + 1;
}
void dfs2(int x, int fa){
	if(son[x]) f[son[x]] = f[x] + 1, g[son[x]] = g[x] - 1, dfs2(son[x], x);
	f[x][0] = 1, ans += g[x][0];
	for(int i = Head[x]; i; i = Next[i]){
		int y = to[i]; if(y == fa || y == son[x]) continue;
		f[y] = cur, cur += mxd[y] << 1, g[y] = cur, cur += mxd[y] << 1;
		dfs2(y, x);
		//f[x], g[x] 表示是之前子树节点 f 的总和 
		for(int j = 0; j < mxd[y]; ++j){
			if(j) ans += f[x][j - 1] * g[y][j];
			ans += g[x][j + 1] * f[y][j]; 
		}
		for(int j = 0; j < mxd[y]; ++j){
			g[x][j + 1] += f[x][j + 1] * f[y][j];  
			if(j) g[x][j - 1] += g[y][j]; //前缀和 
			f[x][j + 1] += f[y][j]; 
		}
	}
}
signed main(){
	n = read();
	for(int i = 1; i < n; ++i){
		int u = read(), v = read();
		add(u, v), add(v, u);
	}
	dfs1(1, 0), f[1] = cur, cur += mxd[1] << 1, g[1] = cur, cur += mxd[1] << 1;
	dfs2(1, 0);
	printf("%lld\n", ans);
	return 0;
}