834. Sum of Distances in Tree [Hard]

最新推荐文章于 2021-07-27 01:59:07 发布

cainiao_23333

最新推荐文章于 2021-07-27 01:59:07 发布

阅读量133

点赞数

分类专栏： leetcode 文章标签： tree graph double-dfs

本文链接：https://blog.csdn.net/weixin_41876385/article/details/108743955

版权

leetcode 专栏收录该内容

159 篇文章 0 订阅

订阅专栏

https://leetcode.com/problems/sum-of-distances-in-tree/

An undirected, connected tree with N nodes labelled 0...N-1 and N-1 edges are given.

The ith edge connects nodes edges[i][0] and edges[i][1] together.

Return a list ans, where ans[i] is the sum of the distances between node i and all other nodes.

Example 1:

Input: N = 6, edges = [[0,1],[0,2],[2,3],[2,4],[2,5]]
Output: [8,12,6,10,10,10]
Explanation: 
Here is a diagram of the given tree:
  0
 / \
1   2
   /|\
  3 4 5
We can see that dist(0,1) + dist(0,2) + dist(0,3) + dist(0,4) + dist(0,5)
equals 1 + 1 + 2 + 2 + 2 = 8.  Hence, answer[0] = 8, and so on.

Note: 1 <= N <= 10000

算法思路：

相当巧妙，自己也知道有大量重复计算，辅助数组也想到了，但不知如何利用edges进行计算，特殊例子可以到处，但没有推出具体的一般的迭代公式。

难点：

地图的转换，将边的信息转换成广义邻接表
对于父子关系，通过parent进行控制，避免parent -> cur -> parent 环的形成，也避免实际建树操作
相邻点的答案依赖关系，先从任意一个点开始，比如0，求出它的答案，然后求其child答案，然后继续向下dfs
依赖关系的推到借助于count数组，可以画一个一般情况，然后推出公示

solution

Tree diagram illustrating recurrence for ans[child]

Then, as illustrated in the diagram, the answer for xx in the entire tree, is the answer of xx on XX "x@X", plus the answer of yy on YY "y@Y", plus the number of nodes in YY "#(Y)". The last part "#(Y)" is specifically because for any node z in Y, dist(x, z) = dist(y, z) + 1.

By similar reasoning, the answer for yy in the entire tree is ans[y] = x@X + y@Y + #(X). Hence, for neighboring nodes xx and yy, ans[x] - ans[y] = #(Y) - #(X), so if we konw ans[y], ans[x] = ans[y] + #(Y) - #(X).

Tree diagram illustrating recurrence for ans[child]

See this diagram above, in dfs we can just transfer from parent to child，so count[child] is the count number of its and its subtrees，similarly count[parent] is the count number of its and its subtrees，but as we can see above, count[parent] include count[child]，so we cannot just use count[parent] to represent right count of numbers, use N - count[child] to solve it(it is great!).

Here is the c++ code.

class Solution {
public:
    vector<int> sumOfDistancesInTree(int N, vector<vector<int>>& edges) {
        vector<vector<int>> g(N);
        for (int i = 0, len = edges.size(); i < len; i++) {
            g[edges[i][0]].push_back(edges[i][1]);
            g[edges[i][1]].push_back(edges[i][0]);
        }
        // 包含自己和子树的节点数，初始只包含自己，记为1
        vector<int> count(N, 1);
        vector<int> res(N);
        // 计算count数组和res[0]答案
        dfs(0, -1, count, g, res);
        // 重新更新计算其他答案
        dfs2(0, -1, count, g, N, res);
        return res;
    }
private:
    void dfs(int cur, int parent, vector<int>& count, 
             vector<vector<int>>& g, vector<int>& res) {
        for (int child : g[cur]) {
            if (child != parent) {
                dfs(child, cur, count, g, res);
                count[cur] += count[child];
                res[cur] += res[child] + count[child];
            }
            
        }
    }
    
    void dfs2(int cur, int parent, vector<int>& count, 
              vector<vector<int>>& g, const int& N, vector<int>& res) {
        for (int child : g[cur]) {
            if (child != parent) {
                // 基于cur，外部比child多的节点数，就是多的距离值
                res[child] = res[cur] + (N - count[child]) - count[child]; 
                dfs2(child, cur, count, g, N, res);
            }
        }
    } 
};