CF1824B1 LuoTianyi and the Floating Islands

CF1824B1 LuoTianyi and the Floating Islands 题解

题目大意

给定一棵 $n$ 个结点的树，现在有 $k(k\le \min(n,3))$ 个结点上有人。

一个结点是好的当且仅当这个点到所有人的距离之和最小。

求在这 $n$ 个点中随机取 $k$ 个点时，好的结点的期望个数，对 $10^9+7$ 取模。

思路

easy 版本的，首先发现，如果只有 $1$ 个人，那么好的节点有且仅有这个人所在点，期望为 $1$。

如果只有 $3$ 个人，如果三个人在同一条链上，好的节点有且仅有中间那个人所在的点，如果不在同一条链上，那么它形成一个“人”字形，好的节点有且仅有“人”字形的交叉点。期望为 $1$。

如果有 $2$ 个人，那么显然好的节点是这两个人所在点为端点组成的链上的所有点。如果枚举两个端点，时间复杂度过高，很难处理。不妨考虑每个点对答案的贡献。

首先考虑每个点作为 lca 的贡献 $v_i$。在 $n$ 个点中任意选择两个点都会有一个 lca，并会为该点增加一个贡献，因此 $\sum_{i=1}^n{v_i}=n \times (n-1) \div 2$。

考虑剩下的情况，每个点作为非 lca 的情况，当且仅当链的两个端点一个在该点的子树内，另一个不在。因此 $w_i={sz}_i \times (n-{sz}_i)$。

一次 dfs 即可解决问题。

代码

int n, k, ans, sum;
bool vis[maxN + 5];
struct node {
    vector<int> to;
    int sz;
} g[maxN + 5];
int dfs(int x) {
    vis[x] = 1;
    g[x].sz = 1;
    for (auto it : g[x].to) {
        if (!vis[it])
            g[x].sz += dfs(it);
    }
    ans = (ans + 1ll * g[x].sz * (n - g[x].sz) % mod) % mod;
    return g[x].sz;
}
int q_pow(int a, int x, int mod) {
    ll base = a, res = 1;
    while (x) {
        if (x & 1)
            res = res * base % mod;
        base = base * base % mod;
        x >>= 1;
    }
    return res;
}
int inv(int x, int mod) { return q_pow(x, mod - 2, mod); }
int main() {
    ios;
    cin >> n >> k;
    if (k == 1 || k == 3) {
        cout << 1 << endl;
        return 0;
    }
    ans = sum = (1ll * n * (n - 1) / 2) % mod;
    rep(i, 1, n - 1) {
        int u, v;
        cin >> u >> v;
        g[u].to.pb(v);
        g[v].to.pb(u);
    }
    dfs(1);
    ans = 1ll * ans * inv(sum, mod) % mod;
    cout << ans << endl;
    return 0;
}