BZOJ4543 Hotel加强版（长链剖分）

题意

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

题解

转自 CaptainHarryChen 的博客

CODE

代码中的f,gf,gf,g对应题解中的num,waynum,waynum,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);

}