BZOJ4543 Hotel加强版（长链剖分）

题意

求一棵树上，两两距离相等的三个点的三元组（无序）的个数。

题解

转自 CaptainHarryChen 的博客

CODE

代码中的$f,g$对应题解中的$num,way$

#include <bits/stdc++.h>
using namespace std;
typedef long long LL;
inline void rd(int &x) {
	char ch; int flg=1; while(!isdigit(ch=getchar()))if(ch=='-')flg=-flg;
	x = 0; do x=x*10+ch-'0'; while(isdigit(ch=getchar())); x*=flg;
}
const int MAXN = 300005;
int n, m, fir[MAXN], to[MAXN<<1], nxt[MAXN<<1], cnt;
inline void link(int u, int v) {
	to[++cnt] = v, nxt[cnt] = fir[u], fir[u] = cnt;
	to[++cnt] = u, nxt[cnt] = fir[v], fir[v] = cnt;
}
int mxd[MAXN], son[MAXN];
void dfs1(int u, int ff) {
	for(int i = fir[u], v; i; i = nxt[i])
		if((v=to[i]) != ff) {
			dfs1(v, u);
			if(mxd[v] > mxd[son[u]]) son[u] = v;
		}
	mxd[u] = mxd[son[u]] + 1;
}

LL tmp[MAXN<<2], *f[MAXN], *g[MAXN], *id=tmp; LL ans;
void dfs(int u, int ff) {
	if(son[u]) f[son[u]] = f[u] + 1, g[son[u]] = g[u] - 1, dfs(son[u], u);
	f[u][0] = 1; ans += g[u][0];
	for(int i = fir[u], v; i; i = nxt[i])
		if((v=to[i]) != ff && v != son[u]) {
			f[v] = id; id += mxd[v]<<1;
			g[v] = id; id += mxd[v]<<1;
			dfs(v, u);
			for(int j = 0; j < mxd[v]; ++j) {
				if(j) ans += f[u][j-1] * g[v][j];
				ans += g[u][j+1] * f[v][j];
			}
			for(int j = 0; j < mxd[v]; ++j) {
				g[u][j+1] += f[u][j+1] * f[v][j];
				if(j) g[u][j-1] += g[v][j];
				f[u][j+1] += f[v][j];
			}
		}
}
int main () {
	rd(n);
	for(int i = 1, x, y; i < n; ++i)
		rd(x), rd(y), link(x, y);
	dfs1(1, 0);
	f[1] = id; id += mxd[1]<<1;
	g[1] = id; id += mxd[1]<<1;
	dfs(1, 0);
	printf("%lld\n", ans);
}