834.树中距离之和

题目描述

题解

树形DP

先考虑本题的一个简化情况, 求出一棵树中每个节点到其所有子节点的距离之和.

对于以u为根的子树, 它到树中所有节点的距离之和用dp[u]表示. 状态转移方程为:
$$
dp[u] = \sum_{v\in son[u]} dp[v] + sz[v]
$$

其中$son[u]$表示u节点的所有儿子节点. 采用自底向上的思路, 所以要用后续DFS, 经过一遍DFS后dp数组就保留了每个节点到以它为根的子树中所有节点的距离和.

在此基础上来分析本题的思路.

在计算每个节点到其他所有节点的距离和时, 要将该节点作为整棵树的根节点, 然后就可以利用上述的dp数组. 所以问题是如何将该节点变成根节点. 首先将节点v从u中剥离
$$
dp[u] = dp[u] - \left(dp[v] + sz[v]\right)
$$
然后将u加到v上面
$$
dp[v] = dp[v] + \left(dp[u] + sz[u]\right)
$$

public class lc834 {

    int[] dp;
    int[] sz;
    int[] ans;
    List<List<Integer>> graph;

    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<>());
        }

        for (int[] edge : edges) {
            int k = edge[0];
            int v = edge[1];
            graph.get(k).add(v);
            graph.get(v).add(k);
        }

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

    }

    private void dfs2(int u, int f) {
        ans[u] = dp[u];

        for (int v : graph.get(u)) {
            if (v == f){
                continue;
            }

            int pu = dp[u], pv = dp[v];
            int su = sz[u], 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] = pu;
            dp[v] = pv;
            sz[u] = su;
            sz[v] = sv;
        }
    }

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

        for (Integer v : graph.get(u)) {
            if (v == f){
                continue;
            }
            dfs(v, u);
            dp[u] += dp[v] + sz[v];
            sz[u] += sz[v];
        }
    }

}