51Nod - 1405 树的距离之和（树形DP）

1405 树的距离之和

题意

给定一棵无根树，假设它有n个节点，节点编号从1到n，求任意两点之间的距离（最短路径）之和。

分析

树形DP。

首先我们让 \(1\) 为根。要开两个数组 \(up \ down\) 分别记录上面点、下面的点到当前点的距离之和。那么对于每个点答案就是 \(up[i] + down[i]\) 。

\(sons[u]\) 数组表示 \(u\) 以及它下面的所有子孙的数量。

显然 \(down[u]\) 是很好求的，当我们计算到某一点 \(u\) 时，当它的以 v 节点为根的子树递归结束后，有 \(down[u] = down[v] + sons[v]\) ，可以把 \(sons[v]\) 当做下面所有点到 \(u\) 这一点有多少条路径，对于 \(u - v\) 这条边，每一条路径都会算一次贡献。

然后在开个 \(DFS\) 去求 \(up[v]\) ，设 \(u\) 为 \(v\) 的父亲节点，有 \(up[v] = up[u] + (n - sons[u]) + (sons[u] - sons[v]) + (down[u] - down[v] - sons[v])\) ，和上面类似 ，第一个括号算的是所有 u 上面的的节点的数量，第二个括号算的是除了 \(v\) 这棵子树，\(u\) 的其它子树的节点数量，意义就和上面的 \(sons[v]\) 一样，最后一个括号算的是 \(u\) 的其它子树上的节点到 \(u\) 的距离之和。

附上一组数据，模拟完就懂了（树形DP真是在树上找规律啊.....）

7
1 2
2 3
2 4
4 6
4 7
2 5
----
13
8
13
9
13
14
14

code

#include<cstdio>
#include<cstring>
#include<algorithm>
#include<cmath>
#include<set>
#include<vector>
#include<iostream>
#include<map>
using namespace std;
typedef long long ll;
const int MAXN = 1e5 + 10;
const int INF = 1e9;
ll up[MAXN], down[MAXN];
int n, sons[MAXN];
int head[MAXN << 1];
struct edge {
    int to, next;
}e[MAXN << 1];
int cnt = 0;
void addedge(int u, int v) {
    e[cnt].to = v;
    e[cnt].next = head[u];
    head[u] = cnt++;
}
void dfs1(int fa, int u) {
    sons[u] = 1;
    for(int i = head[u]; ~i; i = e[i].next) {
        int v = e[i].to;
        if(v != fa) {
            dfs1(u, v);
            sons[u] += sons[v];
            down[u] += down[v] + sons[v];
        }
    }
}
void dfs2(int fa, int u) {
    for(int i = head[u]; ~i; i = e[i].next) {
        int v = e[i].to;
        if(v != fa) {
            up[v] = up[u] + (n - sons[u]) + (sons[u] - sons[v]) + (down[u] - down[v] - sons[v]);
            dfs2(u, v);
        }
    }
}
int main() {
    scanf("%d", &n);
    memset(head, -1, sizeof head);
    for(int i = 1; i < n; i++) {
        int u, v;
        scanf("%d%d", &u, &v);
        addedge(u, v);
        addedge(v, u);
    }
    dfs1(0, 1);
    dfs2(0, 1);
    for(int i = 1; i <= n; i++) {
        printf("%lld\n", up[i] + down[i]);
    }
    return 0;
}