P3565 [POI2014] HOT-Hotels

OMG_NOIP2024-09-09 16:28

~~~~~ P3565 $POI2014$ HOT-Hotels ~~~~~ 总题单链接

思路

~~~~~ 设 g $u$ $i$ g $u$ $i$ g $u$ $i$ 表示在 u u u 的子树内，距离 u u u 为 i i i 的点的个数。

~~~~~ 设 d p $u$ $i$ dp $u$ $i$ dp $u$ $i$ 表示： u u u 的子树内存在两个点 x , y x,y x,y，设 d i s ( x , l c a ) = d i s ( y , l c a ) = d dis(x,lca)=dis(y,lca)=d dis(x,lca)=dis(y,lca)=d， d i s ( l c a , u ) = k dis(lca,u)=k dis(lca,u)=k， i = d − k i=d-k i=d−k。举个栗子：

~~~~~ 上图中 d p $1$ $1$ = 3 dp $1$ $1$ =3 dp $1$ $1$ =3（{x=4,y=5},{x=4,y=8},{x=6,y=8}）。

~~~~~ 对于每个 u u u：

~~~~~ a n s = a n s + d p $u$ $0$ ans=ans+dp $u$ $0$ ans=ans+dp $u$ $0$

~~~~~ a n s = a n s + ∑ x , y ∈ s o n ( u ) , x ! = y d p $x$ $j + 1$ ∗ g $y$ $j - 1$ ans=ans+\sum_{x,y\in son(u),x!=y}dp $x$ $j+1$ *g $y$ $j-1$ ans=ans+∑x,y∈son(u),x!=ydp $x$ $j+1$ ∗g $y$ $j-1$ ，为什么是 j + 1 j+1 j+1 和 j − 1 j-1 j−1？因为 u u u 和 y y y 已经补了两个，不懂的同学可以画个图看一下。

~~~~~ d p $u$ $i$ = d p $u$ $i$ + g $x$ $i - 1$ ∗ g $y$ $i - 1$ dp $u$ $i$ =dp $u$ $i$ +g $x$ $i-1$ *g $y$ $i-1$ dp $u$ $i$ =dp $u$ $i$ +g $x$ $i-1$ ∗g $y$ $i-1$ ，这是 k = 0 k=0 k=0 的情况。

~~~~~ d p $u$ $i$ = d p $u$ $i$ + d p $v$ $i - 1$ dp $u$ $i$ =dp $u$ $i$ +dp $v$ $i-1$ dp $u$ $i$ =dp $u$ $i$ +dp $v$ $i-1$

~~~~~ 以上公式可以用前缀和做到 O ( N ) O(N) O(N) 转移。

~~~~~ 时间复杂度 O ( N 2 ) O(N^2) O(N2)，空间复杂度 O ( N 2 ) O(N^2) O(N2)。

~~~~~ 发现这道题可以用长链剖分将时间复杂度优化至 O ( N ) O(N) O(N)，但这个以后再将。

代码

cpp 复制代码

#include<bits/stdc++.h>
#define ll long long
using namespace std;
int n;ll ans;vector<int>eg[5002];
int f[5002][5002],g[5002][5002];
inline void dfs(int fa,int p)
{
	f[p][0]=1;
	for(int v:eg[p]) {
		if(v==fa)continue;
		dfs(p,v);
		for(int i=0;i<=n;i++)ans+=g[p][i]*(i==0?0:f[v][i-1])+g[v][i+1]*f[p][i];
		for(int i=0;i<=n;i++)g[p][i]+=f[p][i]*(i==0?0:f[v][i-1])+g[v][i+1];
		for(int i=1;i<=n;i++)f[p][i]+=f[v][i-1];
	}
}
signed main(){
	
	cin>>n;
	for(int i=1;i<n;i++) {
		int u,v;cin>>u>>v;
		eg[u].push_back(v);
		eg[v].push_back(u);
	}
	
	dfs(0,1);
	cout<<ans;
	
	return 0;
}