困难

题目描述

给定一个无向、连通的树。树中有 N 个标记为 0…N-1 的节点以及 N-1 条边 。
第 i 条边连接节点 edges[i][0]和edges[i][1]。
返回一个表示节点 i 与其他所有节点距离之和的列表 ans。

来源，leetcode 每日一题 834. 树中距离之和

例如：

  1. 输入: N = 6, edges = [[0,1],[0,2],[2,3],[2,4],[2,5]]
  2. 输出: [8,12,6,10,10,10]
  3. 解释: 
  4. 如下为给定的树的示意图：
  5.   0
  6.  / \
  7. 1   2
  8.    /|\
  9.   3 4 5
  11. 我们可以计算出 dist(0,1) + dist(0,2) + dist(0,3) + dist(0,4) + dist(0,5) 
  12. 也就是 1 + 1 + 2 + 2 + 2 = 8。 因此，answer[0] = 8，以此类推。

解题思路

  1. 暴力解法，直接将所有节点都当作是根节点，计算一次到他们距离的和。O(N^2)的时间和空间复杂度。

  2. 观察，对于上述例子来说，当获得了以0为根节点的树的距离和之后，我们获得了如下信息

    1. 0(8)
    2. /     \
    3. 1(0)   2(3)
    4. /    |   \
    5. 3(0) 4(0) 5(0)

    当我们要计算1的距离和的时候，将1当作根节点，很容易可以得到它的子树的距离和，即0除了1以外，到其余节点的距离和，只需要减去1对0的贡献就可以了，1对0的贡献表现在两方面，一个是通过1到0的距离增加，这部分的值为 1为根节点的子树距离和+1的子节点个数 (这是因为通过1的对0做贡献的，距离都要加1)。还有一个方面是1本身节点对0的贡献。减去这两部分，就可以得到0作为1的子树根节点的距离和了。然后1的部分计算如下： 1的真实子树的距离和+0作为子树的距离和+0的子节点个数+1 。以此类推。因此，在计算的时候，需要实现两方面的东西，1得到以0作为根节点的距离和，同时保存其各个子节点为子树根节点的距离和，2进行上述的根节点转换操作。

    代码

    1. class Solution {
    2. public:
    3.  vector<int> ans, sz, dp;
    4.  vector<vector<int>> graph;
    6.  void dfs(int u, int f) {
    7.      sz[u] = 1; // 计算子节点个数
    8.      dp[u] = 0;
    9.      for (auto& v: graph[u]) {
    10.          // 防止重复计算进入死循环
    11.          if (v == f) {
    12.              continue;
    13.          }
    14.          dfs(v, u);
    15.          dp[u] += dp[v] + sz[v];
    16.          sz[u] += sz[v];
    17.      }
    18.  }
    19.  // 依次将每个节点都替换成根节点
    20.  void dfs2(int u, int f) {
    21.      ans[u] = dp[u];
    22.      for (auto& v: graph[u]) {
    23.          if (v == f) {
    24.              continue;
    25.          }
    26.          int pu = dp[u], pv = dp[v];
    27.          int su = sz[u], sv = sz[v];
    29.          dp[u] -= dp[v] + sz[v];
    30.          sz[u] -= sz[v];
    31.          dp[v] += dp[u] + sz[u];
    32.          sz[v] += sz[u];
    34.          dfs2(v, u);
    36.          dp[u] = pu, dp[v] = pv;
    37.          sz[u] = su, sz[v] = sv;
    38.      }
    39.  }
    41.  vector<int> sumOfDistancesInTree(int N, vector<vector<int>>& edges) {
    42.      ans.resize(N, 0);
    43.      sz.resize(N, 0);
    44.      dp.resize(N, 0);
    45.      graph.resize(N, {});
    46.      // 构建图/树
    47.      for (auto& edge: edges) {
    48.          int u = edge[0], v = edge[1];
    49.          graph[u].emplace_back(v);
    50.          graph[v].emplace_back(u);
    51.      }
    52.      dfs(0, -1);
    53.      dfs2(0, -1);
    54.      return ans;
    55.  }
    56. };