牛客练习赛55 E.树 树形dp

原题链接：https://ac.nowcoder.com/acm/contest/2927/E

题意

有一棵树，设$dis(x,y)$表示x到y的距离，求

$$\sum _{i=1}^n\sum _{j=1}^{n}di{s}^2(i,j)$$

分析

树上路径有一种计算方法

$$dis(x,y)=dep[x]+dep[y]-2\ast dep[lca(x,y)]$$

那么

$$di{s}^2(x,y)=(dep[x]+dep[y]-2\ast dep[lca(x,y)]{)}^2$$

$$=dep[x{]}^2+dep[y{]}^2+2\ast dep[x]dep[y]-4\ast dep[lca(x,y)](dep[x]+dep[y])+4\ast dep[lca(x,y){]}^2$$

对于式子的前三项，可以在处理完所有深度之后O（n）计算得到，接着就是怎么算后面的部分

我们假设lca(x,y)=a，然后对于每一个a，我们对树进行遍历。对于a节点来说，假设b1,b2…bm是a节点的儿子节点，那么点对个数num可以轻松得到
$$num=siz[a]\ast siz[a]-\sum _{i=1}siz[b_i]\ast siz[b_i]$$

易得$\sum 4\ast dep[a{]}^2=4\ast dep[a{]}^2\ast num$

然后考虑$\sum 4\ast dep[a](dep[x]+dep[y])$，首先突破口就是所有点对的深度和，我们可以把深度和转化一下，$dep[x]+dep[y]=2\ast dep[a]+dis(x,y)$，其中dis(x,y)就可以从每个子树开始出发，到除子树外所有点的距离和即$res=\sum (siz[a]-b_i)\ast sum[b_i]$，sum是所有子树节点到该节点的距离之和。

可以推出$\sum 4\ast dep[a](dep[x]+dep[y])=4\ast res\ast dep[a]$

Code

#include <bits/stdc++.h>
using namespace std;
#define fi first
#define se second
#define re register
typedef long long ll;
typedef pair<ll, ll> PII;
typedef unsigned long long ull;
const int N = 1e6 + 10, M = 1e6 + 5, INF = 0x3f3f3f3f;
const int MOD = 998244353;
vector<int> g[N];
ll siz[N], dep[N], sum[N], ans;
int n;
void dfs1(int x, int fa) {
    siz[x] = 1, dep[x] = dep[fa] + 1, sum[x] = dep[x];
    for (auto v : g[x]) {
        if (v == fa) continue;
        dfs1(v, x);
        siz[x] += siz[v];
        sum[x] += sum[v];
    }
}
void dfs2(int x, int fa) {
    ll num = siz[x] * siz[x] % MOD;
    ll res = 2 * dep[x] * siz[x] % MOD;
    for (auto v : g[x]) {
        if (v == fa) continue;
        dfs2(v, x);
        num = (num - siz[v] * siz[v] % MOD + MOD) % MOD;
        res = (res + 2 * sum[v] * (siz[x] - siz[v]) % MOD) % MOD;
    }
    ans = (ans + 4 * dep[x] % MOD * dep[x] % MOD * num % MOD - 4 * dep[x] % MOD * res % MOD + MOD) % MOD;
}
void solve() {
    cin >> n;
    for (int i = 1; i <= n-1; i ++) {
        int u, v; cin >> u >> v;
        g[u].push_back(v);
        g[v].push_back(u);
    }
    dfs1(1, 0);
    dfs2(1, 0);
    ll tmp = 0;
    for (int i = 1; i <= n; i++) ans = (ans + 2 * dep[i] * dep[i] % MOD * n % MOD) % MOD, tmp = (tmp + dep[i]) % MOD;
    ans = (ans + 2 * tmp * tmp % MOD) % MOD;
    cout << ans << endl;
}

signed main() {
    ios::sync_with_stdio(0);
    cin.tie(0);
    cout.tie(0);
#ifdef ACM_LOCAL
    freopen("input", "r", stdin);
    freopen("output", "w", stdout);
#endif
    solve();
}