【Leetcode算法】Sum of Distances in Tree 树中距离之和

Sum of Distances in Tree 树中距离之和

问题描述：

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

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

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

$$\begin{gathered} 1<=n<=3\ast 10^4 \\ edges.length==n-1 \\ edges[i].length==2 \\ 0<=a_i,b_i<n \\ {a}_i!=b_i \\ 给定的输入保证为有效的树 \end{gathered}$$

分析

这是一个以边数组的格式给出的树。要计算出每个节点到其他的节点的距离之和。

对于单个节点，一定是使用dfs，依次计算出每个节点距离root的depth，然后累加，时间复杂度$O(N)$

但是要每一个都跑一遍dfs，很明显是$O(N^2)$.

所以需要其他的方法，来更有效的计算。

策略

如果以a 为root，计算得到的distance为x，下一个计算a的一个子节点b，实际上是不需要重新计算所有的节点距离，只需要知道距离哪些节点近了，距离哪些节点远了，那么以b在内的子树所有节点数量$child[b]$，到达b的距离相较于之前到达a的距离会减少1,而非b的子树节点数量$n-child[b]$到达b的距离会都增加1。

所以此时b到其他所有节点的距离就是$distance[a]+n-child[b]-child[b].$

因此首先就是通过边构建出树，然后计算出每个节点的子树节点数量，最后计算出每个节点的distance。

代码

public List<Integer>[] g;
    public int n;
    int[] ans,arr;
    public int[] sumOfDistancesInTree(int n, int[][] edges) {
        this.n = n;
        ans = new int[n];
        arr = new int[n];
        g = new ArrayList[n];
        for(int i =0;i<n;i++){
            g[i] =new ArrayList();
        }
        for(int[] e: edges){
            int x = e[0],y = e[1];
            g[x].add(y);
            g[y].add(x);
        }
        dfs(0,-1,0);
        reroot(0,-1);
        return ans;
    }

    public int dfs(int node,int fa,int d){
        int sum = 1;
        ans[0]+=d;
        List<Integer> child = g[node];
        for(int c: child){
            if(c==fa)continue;
            sum += dfs(c,node,d+1);
        }
        arr[node] = sum ;
        return sum;
    }
    public void reroot(int cur,int fa){
        List<Integer> child = g[cur];
        for(int c: child){
            if(c==fa)continue;
            ans[c] = ans[cur]+ n- 2*arr[c] ;
            reroot(c,cur );
        }
        return;
    }

时间复杂度$O(N)$

空间复杂度$O(N)$

Tag

Tree

DFS

Graph