首页 > 代码库 > 【DFS】BZOJ3522-[Poi2014]Hotel

【DFS】BZOJ3522-[Poi2014]Hotel

【题目大意】

给出一棵树,求三个节点使得它们两两之间的距离相等,问共有多少种可能性?

【思路】

显然,这三个节点是关于一个中心点对称地辐射出去的。

枚举中心点,往它的各个子树跑Dfs。tmp[i]表示当前子树深度为i的节点个数,p1[i]表示之前的子树中(不包括当前的子树),深度为i的节点的个数,p2[i]表示到当前子树之前,选取两个不同子树的节点的方案数。

对于当前这一棵子树的深度i,显然有ans+=tmp[i]*p2[i]。然后更新下次要用的p1[]、p2[],p2[i]+=p1[i]*tmp[i],p1[i]+=tmp[i]。

不要忘了long long。

 1 #include<bits/stdc++.h>
 2 using namespace std;
 3 typedef long long ll;
 4 const int MAXN=5000+50;
 5 int n;
 6 int p1[MAXN],p2[MAXN],tmp[MAXN];
 7 vector<int> E[MAXN];
 8 
 9 void dfs(int fr,int u,int dep)
10 {
11     tmp[dep]++;
12     for (int i=0;i<E[u].size();i++)
13     {
14         int v=E[u][i];
15         if (v==fr) continue;
16         dfs(u,v,dep+1);
17     }
18 }
19 
20 void solve()
21 {
22     ll ans=0;
23     for (int i=1;i<=n;i++)
24     {
25         memset(p1,0,sizeof(p2));
26         memset(p2,0,sizeof(p1));
27         for (int j=0;j<E[i].size();j++)
28         {
29             memset(tmp,0,sizeof(tmp));
30             int to=E[i][j];
31             dfs(i,to,1);
32             for (int k=1;k<=n;k++)
33             {
34                 ans+=(ll)p2[k]*tmp[k];
35                 p2[k]+=p1[k]*tmp[k];
36                 p1[k]+=tmp[k];
37             }
38         }
39     }
40     printf("%lld\n",ans);
41 }
42 
43 void init()
44 {
45     scanf("%d",&n);
46     for (int i=0;i<n-1;i++)
47     {
48         int x,y;
49         scanf("%d%d",&x,&y);
50         E[x].push_back(y);
51         E[y].push_back(x);
52     }
53 }
54 
55 int main()
56 {
57     init();
58     solve();
59     return 0;    
60 } 

 

【DFS】BZOJ3522-[Poi2014]Hotel