题目链接:HDU 6810
给定一棵 $n$ 个节点的树,两点之间的距离定义为最短路径的边数。定义点集 $S$ 的权值为:
$$ f(S) = \min_{u \in \mathbb{V}} \left(\sum_{v \in S} \text{dist}(u, v)\right) $$
对于所有大小为 $m$ 的点集 $S$,求出 $f(S)$ 的总和,即求 $\sum_{\lvert S\rvert = m} f(S)$。
本题有 $T$ 组数据。
数据范围:$1 \le T \le 1000$,$1 \le m \le n \le 10 ^ {6}$。
Solution
考虑对于每条边计算贡献。
记点集 $S$ 中的点为「关键点」。假设一条边的两侧分别有 $i, m - i$ 个「关键点」,那么这条边至少要被经过 $\min(i, m - i)$ 次,这个下界可以在这棵树的带权重心取到。
因此,设某条边两侧分别有 $s, n - s$ 个节点,那么这条边的贡献为:
$$ f(s) = \sum_{i = 1} ^ {m - 1} \binom{s}{i} \binom{n - s}{m - i} \cdot \min(i, m - i) $$
将 $\min$ 拆成两部分:
$$ \begin{align} f(s) & = \sum_{i = 1} ^ {\left\lfloor\frac{m - 1}{2}\right\rfloor} \binom{s}{i} \binom{n - s}{m - i} \cdot i + \binom{s}{m - i} \binom{n - s}{i} \cdot i \\ & + [m \bmod 2 = 0] \cdot \binom{s}{m / 2} \binom{n - s}{m / 2} \cdot (m / 2) \end{align} $$
由于对称性,我们只需要计算前半部分。记 $p = \left\lfloor\frac{m - 1}{2}\right\rfloor$,则根据吸收律有:
$$ \begin{align} g(s) & = \sum_{i = 1} ^ {p} \binom{s}{i} \binom{n - s}{m - i} \cdot i \\ & = s \sum_{i = 1} ^ {p} \binom{s - 1}{i - 1}\binom{n - s}{m - i} \end{align} $$
我们需要对于所有的 $s = 1 \ldots n - 1$ 计算 $g(s)$。
赋予 $g(s)$ 组合意义:从 $n - 1$ 个球中选出 $m - 1$ 个,其中前 $s - 1$ 个求最多只能选择 $p - 1$ 个。
当 $s = 1$ 时 $g(1) = s \binom{n - 1}{m - 1}$。
当 $s$ 从 $s - 1$ 变为 $s$ 时,不合法的方案为:在前 $s - 2$ 个球中选了 $p - 1$ 个,第 $s - 1$ 个球也被选中,在后 $n - s$ 个球中选择 $m - p - 1$ 个球。
对于 $p = 0$ 的情况需要特殊处理。
时间复杂度:$\mathcal O(n)$。
Code
#include <bits/stdc++.h>
typedef long long int64;
const int N = 1e6, MOD = 1e9 + 7;
int n, m, tot, ans, siz[N + 5], lnk[N + 5], ver[N + 5], nxt[N + 5], fac[N + 5], ifac[N + 5], f[N + 5];
inline int power(int x, int k) {
int ans = 1;
for (; k > 0; k >>= 1, x = (int64)x * x % MOD) {
if (k & 1) {
ans = (int64)ans * x % MOD;
}
}
return ans;
}
inline int inver(int x) {
return power(x, MOD - 2);
}
inline void addEdge(int u, int v) {
ver[++tot] = v;
nxt[tot] = lnk[u];
lnk[u] = tot;
}
inline int binom(int n, int m) {
return m < 0 || n < m ? 0 : (int64)fac[n] * ifac[m] % MOD * ifac[n - m] % MOD;
}
inline int calc(int s) {
int ans = ((int64)s * f[s] % MOD + (int64)(n - s) * f[n - s] % MOD) % MOD;
return (m & 1) ? ans : (ans + (int64)binom(s, m / 2) * binom(n - s, m / 2) % MOD * (m / 2) % MOD) % MOD;
}
void dfs(int u) {
siz[u] = 1;
for (int i = lnk[u]; i > 0; i = nxt[i]) {
int v = ver[i];
dfs(v);
siz[u] += siz[v];
ans = (ans + calc(siz[v])) % MOD;
}
}
int main() {
fac[0] = 1;
for (int i = 1; i <= N; i++) {
fac[i] = (int64)fac[i - 1] * i % MOD;
}
ifac[N] = inver(fac[N]);
for (int i = N; i >= 1; i--) {
ifac[i - 1] = (int64)ifac[i] * i % MOD;
}
int Tcase;
scanf("%d", &Tcase);
for (int cs = 1; cs <= Tcase; cs++) {
scanf("%d%d", &n, &m);
std::fill(lnk + 1, lnk + n + 1, 0);
tot = 0;
for (int p, i = 2; i <= n; i++) {
scanf("%d", &p);
addEdge(p, i);
}
int p = (m - 1) / 2;
if (p == 0) {
std::fill(f + 1, f + n + 1, 0);
} else {
f[1] = binom(n - 1, m - 1);
for (int i = 2; i < n; i++) {
f[i] = (f[i - 1] - (int64)binom(i - 2, p - 1) * binom(n - i, m - p - 1) % MOD + MOD) % MOD;
}
}
ans = 0;
dfs(1);
printf("%d\n", ans);
}
return 0;
}