题解 洛谷 P3565 【[POI2014]HOT-Hotels】
题意:
给定一棵树,在树上选 \(3\) 个点,要求两两距离相等,求方案数。
Solution:
来个题解区没有的做法。
我们可以枚举中间点,枚举中间点的子树,再枚举下距离。
考虑子树的个数:
\(3\) 个:对答案的贡献即为 \(x_1 \times x_2 \times x_3\)
\(4\) 个:对答案贡献为 \(x_1 \times x_2 \times x_3+x_1 \times x_2 \times x_4+x_1 \times x_3 \times x_4+x_2 \times x_3 \times x_4\),即在 \(3\) 个的基础上增加 \(x_4\times (x_1\times x_2+x_1\times x_3+x_2\times x_3)\)。
同理可得 \(5\) 个时,在 \(4\) 个的基础上增加 \(x_5\times(x_1\times x_2+x_1\times x_3+x_1\times x_4+x_2\times x_3+x_2\times x_4+x_3\times x_4)\)。
增加的用 \(sum2\) 维护。
\(x_4\times (x_1\times x_2+x_1\times x_3+x_2\times x_3) \Rightarrow x_5\times(x_1\times x_2+x_1\times x_3+x_1\times x_4+x_2\times x_3+x_2\times x_4+x_3\times x_4)\),增加了 \(x_5\times(x_1+x_2+x_3)\),再用 \(sum1\) 维护即可。
Code:
/*
* @Author: FuTianyu
* @Date: 2020-11-08 18:44:32
* @Last Modified by: FuTianyu
* @Last Modified time: 2020-11-08 19:18:09
*/
#include<bits/stdc++.h>
using namespace std;
#define int long long
#define FOR(i,a,b) for(int i=a;i<=b;i++)
#define REP(i,a,b) for(int i=a;i>=b;i--)
/*
快读快输
*/
const int N=5e3+5;
int n,fir[N],tot;
struct edge{
int nxt,to;
}e[N<<1];
void add(int u,int v){
e[++tot]={fir[u],v};
fir[u]=tot;
e[++tot]={fir[v],u};
fir[v]=tot;
}
int MaxDeep,dep[N],Ans,sum1[N],sum2[N];
void dfs(int u,int dad,int Deep){
MaxDeep=max(MaxDeep,Deep+1);
dep[Deep]++;
for(int i=fir[u];i;i=e[i].nxt){
int v=e[i].to;
if(v==dad) continue;
dfs(v,u,Deep+1);
}
}
void work(int u){
for(int i=fir[u];i;i=e[i].nxt){
int v=e[i].to;
MaxDeep=1;
dfs(v,u,1);
FOR(j,1,MaxDeep){
Ans+=sum2[j]*dep[j];
sum2[j]+=dep[j]*sum1[j];
sum1[j]+=dep[j];
dep[j]=0;
}
}
}
signed main(){
n=read();
FOR(i,1,n-1) add(read(),read());
FOR(i,1,n){
memset(sum1,0,sizeof sum1);
memset(sum2,0,sizeof sum2);
work(i);
}
write(Ans);
return 0;
}