算法题 树中距离之和

834. 树中距离之和

问题描述

给定一个有n个节点的无向树，节点编号从0到n-1。给你一个长度为n-1的二维数组edges，其中edges[i] = [ai, bi]表示树中节点ai和bi之间存在一条边。

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

示例：

输入:n=6,edges=[[0,1],[0,2],[2,3],[2,4],[2,5]]输出:[8,12,6,10,10,10]

• 节点0到其他节点的距离：1→1, 2→1, 3→2, 4→2, 5→2，总和=8
• 节点1到其他节点的距离：0→1, 2→2, 3→3, 4→3, 5→3，总和=12
• 节点2到其他节点的距离：0→1, 1→2, 3→1, 4→1, 5→1，总和=6
• 以此类推

算法思路

两次DFS（树形DP）：

1. 核心：

   • 如果暴力计算每个节点到其他所有节点的距离，时间复杂度为 O(n²)
   • 利用树的结构，可以通过换根DP优化到 O(n)

2. 第一次DFS（后序遍历）：

   • 以任意节点（如0）为根，计算：
     • subtreeSize[node]：以node为根的子树中节点数量
     • dp[node]：子树中所有节点到node的距离之和

3. 第二次DFS（前序遍历）：

   • 利用父节点的结果推导子节点的结果
   • 当从父节点u移动到子节点v时：
     • v的子树中的subtreeSize[v]个节点距离减少1
     • 其他n - subtreeSize[v]个节点距离增加1
     • answer[v] = answer[u] - subtreeSize[v] + (n - subtreeSize[v])
     • answer[v] = answer[u] + n - 2 * subtreeSize[v]

代码实现

importjava.util.*;classSolution{privateList<List<Integer>>graph;// 邻接表表示树privateint[]subtreeSize;// subtreeSize[i] 表示以i为根的子树节点数privateint[]dp;// dp[i] 表示子树中所有节点到i的距离之和privateint[]answer;// 最终答案privateintn;// 节点总数/** * 计算树中每个节点到其他所有节点的距离之和 * 使用两次DFS的换根DP * * @param n 节点数量 * @param edges 边数组 * @return 每个节点的距离和数组 */publicint[]sumOfDistancesInTree(intn,int[][]edges){this.n=n;this.graph=newArrayList<>();this.subtreeSize=newint[n];this.dp=newint[n];this.answer=newint[n];// 初始化邻接表for(inti=0;i<n;i++){graph.add(newArrayList<>());}// 构建无向图for(int[]edge:edges){graph.get(edge[0]).add(edge[1]);graph.get(edge[1]).add(edge[0]);}// 第一次DFS：以0为根，计算子树大小和dp值dfs1(0,-1);// 第二次DFS：换根计算所有节点的答案dfs2(0,-1);returnanswer;}/** * 第一次DFS（后序遍历）： * 计算每个节点的子树大小和子树内距离和 * * @param node 当前节点 * @param parent 父节点（避免回溯） */privatevoiddfs1(intnode,intparent){subtreeSize[node]=1;// 至少包含自己dp[node]=0;// 初始化距离和为0// 遍历所有子节点for(intchild:graph.get(node)){if(child==parent){continue;// 跳过父节点}dfs1(child,node);// 递归处理子树// 累加子树信息subtreeSize[node]+=subtreeSize[child];// 子树中每个节点到当前节点的距离 = 到子节点的距离 + 1// 总距离 = dp[child] + subtreeSize[child]dp[node]+=dp[child]+subtreeSize[child];}}/** * 第二次DFS（前序遍历）： * * @param node 当前节点 * @param parent 父节点 */privatevoiddfs2(intnode,intparent){// 当前节点的答案就是dp[node]（第一次DFS已计算根节点答案）answer[node]=dp[node];// 遍历所有子节点，进行换根for(intchild:graph.get(node)){if(child==parent){continue;// 跳过父节点}// 换根：从node移动到child// 保存原始值intoriginalSubtreeSizeNode=subtreeSize[node];intoriginalSubtreeSizeChild=subtreeSize[child];intoriginalDpNode=dp[node];intoriginalDpChild=dp[child];// 更新子树大小（换根后）subtreeSize[node]=n-subtreeSize[child];subtreeSize[child]=n;// 更新dp值：answer[child] = answer[node] + n - 2 * subtreeSize[child]（换根前的值）dp[child]=dp[node]-subtreeSize[child]+(n-subtreeSize[child]);// 递归处理子树dfs2(child,node);}}}

算法分析

• 时间复杂度：O(n)
  • 两次DFS，每次遍历所有节点和边一次
  • 树有 n-1 条边，所以总时间复杂度为 O(n)
• 空间复杂度：O(n)
  • 邻接表存储：O(n)
  • 递归栈深度：O(n)（最坏情况下树退化为链）

算法过程

输入：n = 6, edges = [[0,1],[0,2],[2,3],[2,4],[2,5]]

第一次DFS（以0为根）：

dfs1(0, -1)：

• 处理子节点1：dfs1(1, 0)
  • subtreeSize[1] = 1,answer[1] = 0
• 处理子节点2：dfs1(2, 0)
  • dfs1(3, 2)→subtreeSize[3] = 1,answer[3] = 0
  • dfs1(4, 2)→subtreeSize[4] = 1,answer[4] = 0
  • dfs1(5, 2)→subtreeSize[5] = 1,answer[5] = 0
  • subtreeSize[2] = 1 + 1 + 1 + 1 = 4
  • answer[2] = (0+1) + (0+1) + (0+1) = 3
• subtreeSize[0] = 1 + 1 + 4 = 6
• answer[0] = (0+1) + (3+4) = 8

第一次DFS后：

• subtreeSize = [6, 1, 4, 1, 1, 1]
• answer = [8, 0, 3, 0, 0, 0]

第二次DFS（换根）：

dfs2(0, -1)：

• 处理子节点1：
  • answer[1] = answer[0] + 6 - 2 * subtreeSize[1] = 8 + 6 - 2 = 12
  • dfs2(1, 0)（1没有其他子节点）
• 处理子节点2：
  • answer[2] = answer[0] + 6 - 2 * subtreeSize[2] = 8 + 6 - 8 = 6
  • dfs2(2, 0)：
    • 处理子节点3：answer[3] = 6 + 6 - 2 = 10
    • 处理子节点4：answer[4] = 6 + 6 - 2 = 10
    • 处理子节点5：answer[5] = 6 + 6 - 2 = 10

最终结果：

answer = [8, 12, 6, 10, 10, 10]

测试用例

importjava.util.Arrays;publicclassMain{publicstaticvoidmain(String[]args){Solutionsolution=newSolution();// 测试用例1：标准示例intn1=6;int[][]edges1={{0,1},{0,2},{2,3},{2,4},{2,5}};int[]result1=solution.sumOfDistancesInTree(n1,edges1);System.out.println("Test 1: "+Arrays.toString(result1));// [8, 12, 6, 10, 10, 10]// 测试用例2：两个节点intn2=2;int[][]edges2={{0,1}};int[]result2=solution.sumOfDistancesInTree(n2,edges2);System.out.println("Test 2: "+Arrays.toString(result2));// [1, 1]// 测试用例3：链状树intn3=4;int[][]edges3={{0,1},{1,2},{2,3}};int[]result3=solution.sumOfDistancesInTree(n3,edges3);System.out.println("Test 3: "+Arrays.toString(result3));// [6, 4, 4, 6]// 测试用例4：星形树intn4=5;int[][]edges4={{0,1},{0,2},{0,3},{0,4}};int[]result4=solution.sumOfDistancesInTree(n4,edges4);System.out.println("Test 4: "+Arrays.toString(result4));// [4, 6, 6, 6, 6]// 测试用例5：单节点intn5=1;int[][]edges5={};int[]result5=solution.sumOfDistancesInTree(n5,edges5);System.out.println("Test 5: "+Arrays.toString(result5));// [0]// 测试用例6：三个节点线性intn6=3;int[][]edges6={{0,1},{1,2}};int[]result6=solution.sumOfDistancesInTree(n6,edges6);System.out.println("Test 6: "+Arrays.toString(result6));// [3, 2, 3]}}

关键点

1. 换根DP：

   • 利用已计算的父节点结果快速推导子节点结果
   • 避免重复计算，将O(n²)优化到O(n)

2. 换根公式：

   • 当根从u移到v时：
     • v的子树中subtreeSize[v]个节点距离-1
     • 其余n - subtreeSize[v]个节点距离+1
     • 变化：-subtreeSize[v] + (n - subtreeSize[v]) = n - 2*subtreeSize[v]

3. 树：

   • 使用邻接表存储无向树
   • DFS时通过parent参数避免回溯

4. 两次遍历：

   • 第一次DFS（后序）：自底向上计算子树信息
   • 第二次DFS（前序）：自顶向下传播换根结果

常见问题

1. 为什么第一次DFS后只有根节点正确？
   第一次DFS只计算了子树内的距离和，没有考虑父节点方向的节点。只有根节点没有父节点。