Problem: 834. 树中距离之和

思路

首先，第一眼就肯定是DP来做

解题方法

先统计一个节点的距离之和，同时统计子树的size！ 本题的关键是从根节点出发（），然后进行换根DP即可

复杂度

• 时间复杂度:

添加时间复杂度, 示例： $O(n)$

• 空间复杂度:

添加空间复杂度, 示例： $O(n)$

Code

class Solution {
    private List<Integer>[] g;
    private int[] ans, size;


// 从根节点出发DFS，累加根节点到每个点的距离，得出ans[0]
// 同时统计每个子树的size
// ans[y] = ans[x] + (n - size[y]) - size[y];
    public int[] sumOfDistancesInTree(int n, int[][] edges) {
        g = new ArrayList[n];
        Arrays.setAll(g, e -> new ArrayList<>());
        for (var e : edges) {
            int x = e[0];
            int y = e[1];
            g[x].add(y);
            g[y].add(x);
        }
        ans = new int[n];
        size = new int[n];
        dfs(0, -1, 0);
        reroot(0, -1);
        return ans;
    }

    private void dfs(int x, int fa, int depth) {
        ans[0] += depth;
        size[x] = 1;
        for (int y : g[x]) {
            if (y != fa) {
                dfs(y, x, depth + 1);
                size[x] += size[y]; // 累加子树size
            }
        }
    }

    private void reroot(int x, int fa) {
        for (int y : g[x]) {
            if (y != fa) {
                ans[y] = ans[x] + g.length - 2 * size[y];
                reroot(y, x);
            }
        }
    }
}