树上DP 代码注释 834.树中距离之和

834.树中距离之和

给定一个无向、连通的树。树中有 N 个标记为 0…N-1 的节点以及 N-1 条边 。
第 i 条边连接节点 edges[i][0] 和 edges[i][1] 。
返回一个表示节点 i 与其他所有节点距离之和的列表 ans。

来源：力扣（LeetCode）
链接：https://leetcode-cn.com/problems/sum-of-distances-in-tree/solution/shu-zhong-ju-chi-zhi-he-by-leetcode-solution/

List<List<Integer>> graph;//图
int[] dp;//dp数组,以各个节点为根到其他节点的距离之和
int[] sz;//以各节点为根的子树（包含节点本身）节点总个数
int[] ans;//最终结果
public int[] sumOfDistancesInTree(int N,int[][] edges){
    // 初始化
    dp = new int[N];
    sz = new int[N];
    ans = new int[N];
    graph = new ArrayList<>();
    for (int i = 0; i < N; i++) {
        graph.add(new ArrayList<Integer>());
    }
    // 创建图
    for (int[] edge : edges){
        int u = edge[0];
        int v = edge[1];
        graph.get(u).add(v);
        graph.get(v).add(u);
    }

    dfs(0,-1);
    dfs2(0,-1);

    return ans;
}

private void dfs(int u, int f) {
    //初始化
    dp[u] = 0;
    sz[u] = 1;

    for (int v : graph.get(u)){//遍历每条边
        if(v == f)// 防止死循环
            continue;

        dfs(v,u);
        dp[u] += dp[v] + sz[v];//
        sz[u] += sz[v];
    }
}

private void dfs2(int u, int f) {
    // 保存最终结果
    ans[u] = dp[u];
    //节点交换
    for (int v : graph.get(u)){
        if (v == f) //防止死循环
            continue;

        int du = dp[u];//保存变量
        int dv = dp[v];
        int su = sz[u];
        int sv = sz[v];

        dp[u] -= (dp[v] + sz[v]);
        sz[u] -= sz[v];
        dp[v] += (dp[u] + sz[u]);
        sz[v] += sz[u];

        dfs2(v,u);

        dp[u] = du;//恢复原状
        dp[v] = dv;
        sz[u] = su;
        sz[v] = sv;
    }
}