【每日一题Day267】LC834树中距离之和 | 换根dp

树中距离之和【LC834】

给定一个无向、连通的树。树中有 n 个标记为 0...n-1 的节点以及 n-1 条边 。

给定整数 n 和数组 edges ， edges[i] = [ai, bi]表示树中的节点 ai 和 bi 之间有一条边。

返回长度为 n 的数组 answer ，其中 answer[i] 是树中第 i 个节点与所有其他节点之间的距离之和。

新知识又++

• 思路

  • 预处理：从节点0出发，进行dfs，求出节点0到每个节点的距离；同时求出以节点$i$为根的子树中的节点个数，并记录在$s[i]$中

  • 换根dp：然后从0出发再次dfs，假设u为根，v为孩子，那么当以$u$为根转移到以$v$为根时，会发生以下变化：

    • 所有在$v$的子树上的节点深度减小了1，那么总深度减少了$s[v]$
    • 所有不在$u$的子树上的节点深度增加了1，那么总深度增加了$n-s[v]$

    状态转移方程为
    $$f_v=f_u-s_v+n-s_v=f_u+n-2s_v$$

• 代码

  class Solution {
      private List<Integer>[] g;
      private int[] ans, size;
  
      public int[] sumOfDistancesInTree(int n, int[][] edges) {
          g = new ArrayList[n]; // g[x] 表示 x 的所有邻居
          Arrays.setAll(g, e -> new ArrayList<>());
          for (var e : edges) {
              int x = e[0], y = e[1];
              g[x].add(y);
              g[y].add(x);
          }
          ans = new int[n];
          size = new int[n];
          dfs(0, -1, 0); // 0 没有父节点
          reroot(0, -1); // 0 没有父节点
          return ans;
      }
  
      private void dfs(int x, int fa, int depth) {
          ans[0] += depth; // depth 为 0 到 x 的距离
          size[x] = 1;
          for (int y : g[x]) { // 遍历 x 的邻居 y
              if (y != fa) { // 避免访问父节点
                  dfs(y, x, depth + 1); // x 是 y 的父节点
                  size[x] += size[y]; // 累加 x 的儿子 y 的子树大小
              }
          }
      }
  
      private void reroot(int x, int fa) {
          for (int y : g[x]) { // 遍历 x 的邻居 y
              if (y != fa) { // 避免访问父节点
                  ans[y] = ans[x] + g.length - 2 * size[y];
                  reroot(y, x); // x 是 y 的父节点
              }
          }
      }
  }
  
  作者：灵茶山艾府
  链接：https://leetcode.cn/problems/sum-of-distances-in-tree/solutions/2345592/tu-jie-yi-zhang-tu-miao-dong-huan-gen-dp-6bgb/
  来源：力扣（LeetCode）
  著作权归作者所有。商业转载请联系作者获得授权，非商业转载请注明出处。
  

  • 复杂度
    • 时间复杂度：$\mathcal{O}(n)$
    • 空间复杂度：$\mathcal{O}(n)$