834. 树中距离之和

今天碰到一道hard的leetcode，题型是二叉树，但是解法很新颖————树形动态规划，之前很少见到，感觉题型比较典型，应该后面的面试啥的会遇到，所以在这里记录一下。来源：力扣（LeetCode）
链接：https://leetcode-cn.com/problems/sum-of-distances-in-tree

题目描述：给定一个无向、连通的树。树中有 N 个标记为 0...N-1 的节点以及 N-1 条边 。

第 i 条边连接节点 edges[i][0] 和 edges[i][1] 。

返回一个表示节点 i 与其他所有节点距离之和的列表 ans。

分析：先开始我看的时候，还以为要利用关于图的解法，但是后面参考了答案之后才发现使用的是树形dp解法。本题依赖的框架是求解根到各个节点的距离之和，在此基础上再加上换根的操作，所以代码中共有两次遍历，第一次遍历构造每个dp,sz的数值，第二次遍历进行换根操作，得到需要的答案。主义的是，sz[u]表示u节点的子节点个数，同时还要加上自身。

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

    // 第一次遍历，构建初始的dp,sz表，获得每个节点的dp,sz值
    /*
     * u:当前节点
     * f:当前节点的父节点
     */
    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, {});
        // auto& edge中加上&，相当于浅拷贝，用原地址上的int，不用每次遍历都创建一个新的int变量
        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;
    }
};

算法的时间复杂度为O(N)，空间复杂度为O(N)