P3565 [POI2014] HOT-Hotels

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

思路

~~~~~      设 $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)$。

~~~~~      发现这道题可以用长链剖分将时间复杂度优化至 $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;
}