题目地址：洛谷：【P3565】[POI2014]HOT-Hotels – 洛谷、BZOJ：Problem 3522. — [Poi2014]Hotel

题目描述

有一个树形结构的宾馆，n个房间，n-1条无向边，每条边的长度相同，任意两个房间可以相互到达。吉丽要给他的三个妹子各开（一个）房（间）。三个妹子住的房间要互不相同（否则要打起来了），为了让吉丽满意，你需要让三个房间两两距离相同。
有多少种方案能让吉丽满意？

输入输出格式

输入格式：
第一行一个数n。
接下来n-1行，每行两个数x,y，表示x和y之间有一条边相连。

输出格式：
让吉丽满意的方案数。

输入输出样例

输入样例#1：

7
1 2
5 7
2 5
2 3
5 6
4 5

输出样例#1：

5

说明

【样例解释】
{1,3,5},{2,4,6},{2,4,7},{2,6,7},{4,6,7}
【数据范围】
n≤5000

题解

参考资料：题解 P3565 【[POI2014]HOT-Hotels】 – – 洛谷博客
很容易有一个思路就是枚举树根点拉成有根树，然后从根的不同子树中任取三个子树，从这些子树的同深度点集中每个子树任取一个，求这样的方案数之和即可。
问题是，从一堆集合中任取三个，再从这三个集合中任取一个元素这样的方案数总和是一个很难求的东西。接下来是一波神仙分析。
当只有3个集合，元素个数分别是a、b、c时，显然上面说的方案数就是abc。
当集合变为4个，新加入的集合元素个数为d时，我们会发现答案会增加d(ab+bc+ac)。
当集合变为5个，新加入的集合元素个数为e时，我们会发现答案会增加e(ab+bc+ac+ad+bd+cd)。
……
注意上面括号内的部分，实际上每一次会增加上一次元素个数×上一次之前的元素个数和这么多，我们把括号内的这个东西维护成sum2，其实际意义是任选2个集合再分别从两个集合中各任选1个的方案数，而维护sum1为所有集合的元素个数和。每次新加入一个集合就可以往答案中加入new×sum2，sum2加入new×sum1，sum1加入new，这样就很好维护了。
于是我们解决了这个数学模型的问题。接下来的就是对这一棵子树的每一个深度集合都做一遍这个操作即可。
总复杂度为O(n^2 \log n)。

代码

// Code by KSkun, 2018/6
#include <cstdio>
#include <cctype>
#include <cstring>

#include <algorithm>
#include <vector>

typedef long long LL;

inline char fgc() {
    static char buf[100000], *p1 = buf, *p2 = buf;
    return p1 == p2 && (p2 = (p1 = buf) + fread(buf, 1, 100000, stdin), p1 == p2) 
        ? EOF : *p1++;
}

inline LL readint() {
    register LL res = 0, neg = 1;
    register char c = fgc();
    while(!isdigit(c)) {
        if(c == '-') neg = -1;
        c = fgc();
    }
    while(isdigit(c)) {
        res = (res << 1) + (res << 3) + c - '0';
        c = fgc();
    }
    return res * neg;
}

const int MAXN = 5005;

int n;

std::vector<int> gra[MAXN];

LL ans, sum1[MAXN], sum2[MAXN];
int cnt[MAXN], mxdep;

inline void dfs(int u, int fa, int dep) {
    cnt[dep]++; mxdep = std::max(mxdep, dep);
    for(int i = 0; i < gra[u].size(); i++) {
        int v = gra[u][i];
        if(v == fa) continue;
        dfs(v, u, dep + 1);
    }
}

int main() {
    n = readint();
    for(int i = 1, a, b; i < n; i++) {
        a = readint(); b = readint();
        gra[a].push_back(b); gra[b].push_back(a);
    }
    for(int i = 1; i <= n; i++) {
        memset(sum1, 0, sizeof(sum1));
        memset(sum2, 0, sizeof(sum2));
        for(int j = 0; j < gra[i].size(); j++) {
            int v = gra[i][j];
            mxdep = 0; dfs(v, i, 1);
            for(int k = 1; k <= mxdep; k++) {
                ans += cnt[k] * sum2[k];
                sum2[k] += sum1[k] * cnt[k];
                sum1[k] += cnt[k];
                cnt[k] = 0;
            }
        }
    }
    printf("%lld", ans);
    return 0;
}