834. 树中距离之和 (opens new window)

1.换根dp

换根DP算法：假设y是x的儿子，ans表示节点x的距离之和，size表示子树y的大小(以0节点作为树根形成树结构)，那么字节点y的距离之和通过公式：ans[y]=ans[x]+n-2*size[y]可以计算得到。

• 第一遍DFS：统计以0节点作为根的路径之和ans[0]，同时根据dfs返回子节点数量计算size数组。(注意每次index节点搜索范围需要从所有节点缩小到邻居节点,否则使用visit还会超时)

  

• 第二遍DFS：每个节点y将子节点x进行换根，并利用父节点y的路径之和来计算子节点x的路径和：

  

算法本质：在计算时考虑交换两个相邻节点的父子关系，从而找出两个节点之间距离之和的变化量，而这个变化量依赖于子树的大小。从而根据父节点的距离和ans[y]来递推计算出ans[x]子节点的距离之和。

class Solution {
    int zeroNode=0;
    int[] res;
    List<Integer>[] list;
    public int[] sumOfDistancesInTree(int n, int[][] edges) {
        res=new int[n];
        int[] size=new int[n];
        list=new ArrayList[n];
        Arrays.setAll(list, e -> new ArrayList<>());
        for(int i=0;i<edges.length;i++){
            list[edges[i][0]].add(edges[i][1]);
            list[edges[i][1]].add(edges[i][0]);
        }
        //初始化size数组，同时计算第一个0号节点的所有路径之和
        dfs(edges,size,1,0,-1);
        res[0]=zeroNode;
        //换根
       dfsNext(size,zeroNode,0,-1);
       return res;
    }
    public int dfs(int[][] edges,int[] size,int currlen,int index,int father){
        int count=1;
        for(Integer i:list[index]){
            if(i!=father){
                zeroNode+=currlen;
                count+=dfs(edges,size,currlen+1,i,index);
            }
        }
        size[index]=count;
        return count;
    }
    public void dfsNext(int[] size,int sumlen,int index,int father){
        int n=size.length;
        for(Integer i:list[index]){
            if(i!=father){
                res[i]=sumlen+n-2*size[i];
                dfsNext(size,res[i],i,index);
            }
        }
    }
}