P3565 [POI2014] HOT-Hotels

      ~~~~~      P3565 [POI2014] HOT-Hotels       ~~~~~      总题单链接

      ~~~~~       2024.9.10：DP方程有问题，已修改，同时更新了长链剖分优化版本。

思路

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

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

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

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

      ~~~~~      $ans=ans+dp[u][0]$

      ~~~~~      $ans=ans+{\sum }_{x,y\in son(u),x!=y}dp[x][j+1]\ast g[y][j-1]$，为什么是 $j+1$ 和 $j-1$？因为 $u$ 和 $y$ 已经补了两个，不懂的同学可以画个图看一下。

      ~~~~~      $dp[u][i]=dp[u][i]+g[x][i-1]\ast g[y][i-1]$，这是 $k=0$ 的情况。

      ~~~~~      $dp[u][i]=dp[u][i]+dp[v][i+1]$

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

      ~~~~~      时间复杂度 $O(N^2)$，空间复杂度 $O(N^2)$，理论可以有 $100$ 分，但我只有 $90$ 分。

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

代码

#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;
}