题面描述 树中距离和
给定一棵 $n$ 个节点的树( $1 \leq n \leq 10000$ )，对于每一个节点，求出其到所有节点距离的和。
只要写过不太裸的数位dp或者不太裸的线段树或者对维护东西有一定理解的话，肯定能想到怎么在 $O\left(n\right)$ 的复杂度内求出某一个点到所有点的距离和。方法是维护每棵子树到子树根的距离和，和每棵子树去除根节点的节点数，然后设法转移即可。
设 $sum[p]$ 表示 $p$ 节点所在的子树中，每个节点到 $p$ 的距离和；
设 $cnt[p]$ 表示 $p$ 节点所在子树大小(不包含根).
则对于 $p$ 的父亲 $fa$ 有:
因此，最朴素的做法是 $O\left(n^2\right)$ 的，即将每一个节点作为根节点走一遍dfs。
现在我们试图通过转移在 $O\left(1\right)$ 次数内求出所有点的距离和。
例如对于下图，对于节点 $3$ ，其答案恰好可以通过节点 $2$ 转移。
1 | 0 |
2 | / \ |
3 | 1 2 |
4 | /|\ |
5 | 3 4 5 |
6 | | |
7 | 6 |
节点 $3$ 的距离和恰好为除去包含节点 $3$ 的子树的节点（即 $[0,1,2,4,5]$ ）到 $3$ 的距离和，加上其子树(即 $[6]$ )到 $3$ 的距离和。
其中，前者在已知节点 $2$ 到全体距离和的情况下，可以转移得到，方法和之前提到的朴素做法思想类似；而后者是朴素做法的计算结果。
现设 $ans[fa]$ 表示树上所有节点到 $fa$ 节点的距离和，则对于其子节点 $p$ 有:
因此，本题可以在两遍dfs的复杂度下完成，总体时间复杂度为 $O\left(n\right)$ 。
1 | class Solution { |
2 | int n; |
3 | vector<int> Edges[10010]; |
4 | int cnt[10010]; |
5 | int sum[10010]; |
6 | vector<int> ans; |
7 | bool vis[10010]; |
8 | |
9 | void dfs(int p){ |
10 | for (auto &e:Edges[p]){ |
11 | if (!vis[e]){ |
12 | vis[e]=true; |
13 | dfs(e); |
14 | sum[p]+=sum[e]+cnt[e]+1; |
15 | cnt[p]+=cnt[e]+1; |
16 | } |
17 | } |
18 | } |
19 | |
20 | void dfs2(int p,int fa){ |
21 | ans[p]=ans[fa]+n-(cnt[p]+1)*2; |
22 | //ans[p]=(ans[fa]-sum[p]-cnt[p]-1)+(n-cnt[p]-1)+sum[p]; |
23 | for (auto e:Edges[p]){ |
24 | if (!vis[e]){ |
25 | vis[e]=true; |
26 | dfs2(e,p); |
27 | } |
28 | } |
29 | } |
30 | |
31 | public: |
32 | vector<int> sumOfDistancesInTree(int N, vector<vector<int>>& x) { |
33 | n=N; |
34 | for(auto &e:x){ |
35 | Edges[e[0]].push_back(e[1]); |
36 | Edges[e[1]].push_back(e[0]); |
37 | } |
38 | vis[0]=true; |
39 | dfs(0); |
40 | ans.resize(n); |
41 | ans[0]=sum[0]; |
42 | memset(vis+1,0,sizeof(bool)*(n-1)); |
43 | for (auto e:Edges[0]){ |
44 | vis[e]=true; |
45 | dfs2(e,0); |
46 | } |
47 | return ans; |
48 | } |
49 | }; |
Comment: 这种题属于非常套路的题目，做多了就很容易想出来