BZOJ3522 [Poi2014]Hotel 【树形dp】

题目链接

BZOJ3522

题解

就是询问每个点来自不同子树离它等距的三个点的个数
数据支持\(O(n^2)\)，可以对每个距离分开做
设\(f[i][j]\)表示\(i\)的子树中到\(i\)距离为\(j\)的点的个数
利用换根法可得到每个点作为根时的值
然后随便容斥一下就是答案

#include<algorithm>
#include<iostream>
#include<cstring>
#include<cstdio>
#include<cmath>
#include<map>
#define Redge(u) for (int k = h[u],to; k; k = ed[k].nxt)
#define REP(i,n) for (int i = 1; i <= (n); i++)
#define mp(a,b) make_pair<int,int>(a,b)
#define cls(s) memset(s,0,sizeof(s))
#define cp pair<int,int>
#define LL long long int
using namespace std;
const int maxn = 5005,maxm = 100005,INF = 1000000000;
inline int read(){
	int out = 0,flag = 1; char c = getchar();
	while (c < 48 || c > 57){if (c == '-') flag = -1; c = getchar();}
	while (c >= 48 && c <= 57){out = (out << 3) + (out << 1) + c - 48; c = getchar();}
	return out * flag;
}
int h[maxn],ne = 1;
struct EDGE{int to,nxt;}ed[maxn << 1];
int siz[maxn][maxn],fa[maxn],g[maxn],n;
LL ans;
inline void build(int u,int v){
	ed[++ne] = (EDGE){v,h[u]}; h[u] = ne;
	ed[++ne] = (EDGE){u,h[v]}; h[v] = ne;
}
void dfs(int u){
	siz[u][0] = 1;
	Redge(u) if ((to = ed[k].to) != fa[u]){
		fa[to] = u; dfs(to);
		for (int k = 1; k <= n; k++)
			siz[u][k] += siz[to][k - 1];
	}
}
void dfs2(int u){
	if (u != 1){
		int v = fa[u];
		g[0] = 1;
		for (int k = 1; k <= n; k++)
			g[k] = siz[v][k] - siz[u][k - 1];
		for (int k = 1; k <= n; k++)
			siz[u][k] += g[k - 1];
	}
	for (int K = 1; K <= n; K++){
		LL s1 = g[K - 1],s2 = s1 * s1,s3 = s2 * s1;
		Redge(u) if ((to = ed[k].to) != fa[u]){
			s1 += siz[to][K - 1];
			s2 += 1ll * siz[to][K - 1] * siz[to][K - 1];
			s3 += 1ll * siz[to][K - 1] * siz[to][K - 1] * siz[to][K - 1];
		}
		ans += (s1 * s1 * s1 - 3ll * (s2 * s1 - s3) - s3) / 6;
	}
	Redge(u) if ((to = ed[k].to) != fa[u])
		dfs2(to);
}
int main(){
	n = read();
	for (int i = 1; i < n; i++) build(read(),read());
	dfs(1); dfs2(1);
	printf("%lld\n",ans);
	return 0;
}