P10842 题解

EuphoricStar · 2024-08-04 22:18:33 · 题解

若每条边 (a, b) 维护一个计数器 c_{(a, b)}，每次枚举 u, v, i 时，把 x \to i 路径上的边的 c 值全部加 1，那么答案就是 \sum\limits_{(a, b) \in E} c_{(a, b)}。

考虑直接计算 c_{(a, b)}。设断开 (a, b) 这条边后形成的两棵子树大小分别为 x, y。那么当 u \to v 的路径不跨过 (a, b) 这条边且 i 点在另外一棵子树内，c_{(a, b)} 会加 1。所以实际上只用计算在一棵子树内选两个点，在另一棵子树选一个点的方案数。所以 c_{(a, b)} = y \cdot \frac{x(x - 1)}{2} + x \cdot \frac{y(y - 1)}{2}。枚举每条边直接求和即可。

时间复杂度 O(n)。

#include <bits/stdc++.h>
#define pb emplace_back
#define fst first
#define scd second
#define mkp make_pair
#define mems(a, x) memset((a), (x), sizeof(a))

using namespace std;
typedef long long ll;
typedef double db;
typedef unsigned long long ull;
typedef long double ldb;
typedef pair<ll, ll> pii;

const int maxn = 200100;
const ll mod = 1000000007;

ll n, sz[maxn], ans;
vector<int> G[maxn];

void dfs(int u, int fa) {
    sz[u] = 1;
    for (int v : G[u]) {
        if (v == fa) {
            continue;
        }
        dfs(v, u);
        sz[u] += sz[v];
        ans = (ans + sz[v] * (sz[v] - 1) / 2 * (n - sz[v]) + (n - sz[v]) * (n - sz[v] - 1) / 2 * sz[v]) % mod;
    }
}

void solve() {
    scanf("%lld", &n);
    for (int i = 1, u, v; i < n; ++i) {
        scanf("%d%d", &u, &v);
        G[u].pb(v);
        G[v].pb(u);
    }
    dfs(1, -1);
    printf("%lld\n", ans);
}

int main() {
    int T = 1;
    // scanf("%d", &T);
    while (T--) {
        solve();
    }
    return 0;
}