牛客练习赛55 E 树

题目链接：

题意：给出n个点，n-1条边求任意两个点的距离平方的和

解法：

f[i]表示这个点的高度

sz[i]表示这个子树的大小

szz[i]表示这个这个子树大小的平方

sum[i]表示这个子树所有点高度的和

两个点i, j的距离dis = f[i] + f[j] - 2 * f[lca(i, j)]

dis的平方 =  f[i] * f[i] + f[j] * f[j] + 2 * f[i] * f[j] * 4 * f[lca(i, j)] * f[lca(i, j)]  - 4 *  (f[i] + f[j]) * f[lca(i, j)]

前面三项直接计算即可

计算第四项就要计算点u为lca的(i, j)对数

以点u为lca的(i, j)对数为sz[u] * sz[u] - sz[v] * sz[v](v为u的所有儿子)

计算第五项就要计算以点u为lca的f[i]和f[j]的和

以点u为lca的f[i]和f[j]的和为2 * sum[v] * (sz[u] - sz[v]) (v为u的所有儿子)（乘2是因为（i, j）,(j, i)都要计算）

最后再加上点u自身的贡献:2 * f[u] * sz[u]

#include <bits/stdc++.h>

using namespace std;

typedef long long ll;

const int mod = ;

const int M = 1e6 + ;

ll ans;

int cnt;

int head[M];

ll sz[M], szz[M], f[M], sum[M];

struct Edge{

    int next, to;

}edge[M * ];

void add_edge(int u, int v) {

    edge[++cnt].next = head[u];

    edge[cnt].to = v;

    head[u] = cnt;

}

void dfs(int u, int fa) {

    sz[u] = ;

    f[u] = f[fa] + ;

    sum[u] = f[u];

    for(int i = head[u]; i; i = edge[i].next) {

        int v = edge[i].to;

        if(v == fa) continue;

        dfs(v, u);

        sz[u] += sz[v];

        sum[u] += sum[v];

    }

    szz[u] = (sz[u] % mod) * (sz[u] % mod) % mod;

}

void dfs1(int u, int fa) {

    ll ans1 = szz[u], ans2 = ( * f[u] % mod) * (sz[u] % mod) % mod;

    for(int i = head[u]; i; i = edge[i].next) {

        int v = edge[i].to;

        if(v == fa) continue;

        dfs1(v, u);

        ans1 = (ans1 - szz[v] + mod) % mod;

        ans2 = (ans2 % mod + ( * sum[v] % mod) * ((sz[u] - sz[v]) % mod) % mod) % mod;

    }

    ans = (ans % mod + (( * f[u] % mod) * (f[u] % mod) % mod) * (ans1 % mod) % mod) % mod;

    ans = (ans % mod - ( * ans2) % mod * (f[u] % mod) % mod + mod ) % mod;

}

int main(){

    int n;

    while(~scanf("%d", &n)) {

        cnt = ;

        ans = ;

        memset(head, , sizeof(head));

        for(int i = ; i <= n - ; i++) {

            int u, v;

            scanf("%d%d", &u, &v);

            add_edge(u, v);

            add_edge(v, u);

        }

        dfs(, );

        ll summ = ;

        for(int i = ; i <= n; i++) {

            summ = (summ % mod + f[i] % mod) % mod;

            ans = (ans % mod + ((f[i] % mod) * (f[i] % mod) % mod ) * (( * n) % mod )% mod) % mod;

        }

        ans = (ans % mod + ( * summ % mod) * (summ % mod) % mod) % mod;

        dfs1(, );

        printf("%lld\n", ans);

    }

    return ;

}