834. 树中距离之和， 后序遍历，自底向下的解决方案！

给定一个无向、连通的树。树中有 N 个标记为 0...N-1 的节点以及 N-1 条边 。第 i 条边连接节点 edges[i][0] 和 edges[i][1]。
返回一个表示节点 i 与其他所有节点距离之和的列表 ans。
输入: N = 6, edges = [[0,1],[0,2],[2,3],[2,4],[2,5]]
输出: [8,12,6,10,10,10]
解释:
如下为给定的树的示意图：

  0
 / 1   2
   /|  3 4 5

我们可以计算出 dist(0,1) + dist(0,2) + dist(0,3) + dist(0,4) + dist(0,5)
也就是 1 + 1 + 2 + 2 + 2 = 8。 因此，answer[0] = 8，以此类推。

思路：
很容易想到一个树形动态规划：定义 dp[u] 表示以 u 为根的子树，它的所有子节点到它的距离之和，同时定义 sz[u] 表示以 u 为根的子树的节点数量，不难得出如下的转移方程：

\(sz[v]\) 表示以\(v\)为根节点的子数有多少个节点，包括\(v\)。
我们遍历整棵树，从叶子节点开始自底向上递推到根节点\(root\) 即能得出最后的答案为 \(dp[root]\)。在描述点与点的拓扑关系的时候，常常使用邻接表来表示，临界表的索引往往表示第几个节点，索引指向的数组，表示与这个节点直接相连的有哪些节点。通过上述特性构造点与点之间的拓扑关系。
假设 \(u\) 的某个后代节点（后代节点与夫节点直接相连）为 \(v\)，如果要算 \(v\) 的答案，本来我们要以 \(v\) 为根来进行一次树形动态规划。但是利用已有的信息，我们可以考虑树的形态做一次改变，让 \(v\) 换到根的位置,\(u\) 变为其孩子节点（也就是二叉树旋转一下），同时维护原有的 \(dp\) 信息。在这一次的转变中，我们观察到除了 \(u\) 和 \(v\) 的 \(dp\) 值，其他节点的 \(dp\) 值都不会改变，因此只要更新 \(dp[u]\) 和 \(dp[v]\) 的值即可。

那么我们来看 vv 换到根的位置的时候怎么利用已有信息求出 \textit{dp}[u]dp[u] 和 \textit{dp}[v]dp[v] 的值。重新回顾第一次树形动态规划的转移方程，我们可以知道当 uu 变为 vv 的孩子的时候 vv 不在 uu 的后代集合 \textit{son}[u]son[u] 中了，因此此时 \textit{dp}[u]dp[u] 需要减去 vv 的贡献，即

\textit{dp}[u]=\textit{dp}[u]-(\textit{dp}[v]+\textit{sz}[v])
dp[u]=dp[u]?(dp[v]+sz[v])

同时 \textit{sz}[u]sz[u] 也要相应减去 \textit{sz}[v]sz[v]。

而 vv 的后代节点集合中多出了 uu，因此 \textit{dp}[v]dp[v] 的值要由 uu 更新上来，即

\textit{dp}[v]=\textit{dp}[v]+(\textit{dp}[u]+\textit{sz}[u])
dp[v]=dp[v]+(dp[u]+sz[u])

同时 \textit{sz}[v]sz[v] 也要相应加上 \textit{sz}[u]sz[u]。

class Solution {
public:
    vector<int> ans, sz, dp;
    vector<vector<int>> graph;

    void dfs(int u, int f) {
        sz[u] = 1;
        dp[u] = 0;
        for (auto& v: graph[u]) {
            if (v == f) {
                continue;
            }
            dfs(v, u);
            dp[u] += dp[v] + sz[v];
            sz[u] += sz[v];
        }
    }

    void dfs2(int u, int f) {
        ans[u] = dp[u];
        for (auto& v: graph[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;
        }
    }

    vector<int> sumOfDistancesInTree(int N, vector<vector<int>>& edges) {
        ans.resize(N, 0);
        sz.resize(N, 0);
        dp.resize(N, 0);
        graph.resize(N, {});
        for (auto& edge: edges) {
            int u = edge[0], v = edge[1];
            graph[u].emplace_back(v);
            graph[v].emplace_back(u);
        }
        dfs(0, -1);
        dfs2(0, -1);
        return ans;
    }
};

834. 树中距离之和， 后序遍历，自底向下的解决方案！

原文：https://www.cnblogs.com/wsl-hitsz/p/13773703.html