题目链接: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;
}
最后修改:2020 年 12 月 22 日
如果觉得我的文章对你有用,请随意赞赏