0%

Leetcode834 树上距离和

题面描述 树中距离和

给定一棵 $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: 这种题属于非常套路的题目,做多了就很容易想出来