【Bzoj3522】Hotel

题意

有一个树形结构，每条边的长度相同，任意两个节点可以相互到达。选3个点。两两距离相等。有多少种方案？

---

解析

满足条件的点对一定是“有一个中心点，三个点到中心点的距离相等，且三个点分别在不同子树中”这种情况。(因为如果在同一子树中会被重复统计)。
那么枚举中心点，然后遍历子树，统计答案。
令cnt[i]表示当前子树中深度为i的点有多少个。
f[i]表示当前点已经遍历过的子树中，深度为i的点有多少个。
g[i]表示当前点已经遍历过的子树中，两个点深度都为i的点对有多少个。
根据乘法原理ans += ∑(g[i] * cnt[i])。每次用f和cnt去更新g，用cnt去更新f。

---

#include <cstdio>
#include <cstring>
#include <algorithm>

#define Rep( i , _begin , _end ) for(int i=(_begin),i##_END=(_end);i<=(i##_END);i++)
#define For( i , _begin , _end ) for(int i=(_begin),i##_END=(_end);i!=(i##_END);i++)
#define Lop( i , _begin , _end ) for(int i=(_begin),i##_END=(_end);i>=(i##_END);i--)
#define Dnt( i , _begin , _end ) for(int i=(_begin),i##_END=(_end);i!=(i##_END);i--)

using std :: max;
using std :: min;

typedef long long LL;

const int maxx = 100000 + 25;

int n,m,x,y,z,num;
LL  ans;

int f[maxx],g[maxx],cnt[maxx],son[maxx];
int head[maxx],nxt[maxx],to[maxx];

void Ins(int x,int y){
    to[++num] = y;nxt[num] = head[x];head[x] = num;son[x] ++;
}

void Dfs(int x,int pre,int dpt){
    cnt[dpt] ++;
    for(int i=head[x];i;i=nxt[i])
        if(to[i] != pre)
            Dfs(to[i],x,dpt+1);
}

int main(){
    scanf("%d",&n);
    For( i , 1 , n ) scanf("%d%d",&x,&y),Ins(x,y),Ins(y,x);
    Rep( x , 1 , n ){
        if(son[x] < 3) continue;
        memset(f,0,sizeof(f));
        memset(g,0,sizeof(g));
        for(int i=head[x];i;i=nxt[i]){
            memset(cnt,0,sizeof(cnt));
            Dfs(to[i],x,1);
            Rep( k , 1 , n ){
                ans += (LL)cnt[k] * g[k];
                g[k] += f[k] * cnt[k];
                f[k] += cnt[k];
            }
        }
    }
    printf("%lld\n",ans);
    return 0;
}