树中距离之和（树形dp）

题意：给定一个含有n个节点的树，求各个点到其他点的距离之和。返回一个vector。

举例：假设返回vectot<int> ans。则ans[0] = dist(0,1) + dist(0,2) + ...+ dist(0,n)

思路：考虑一种简化的情况：求各个点到根节点的距离之和。

该简化情况可通过树状dp来解决。

设数组dp[]，dp[u]为以点u为根节点的子树的子节点到点u的距离之和；数组nodes[]，nodes[u]为以点u为根节点的节点数。

则更新公式为：

 

所以，在以root为根节点的树中，其他点到根节点root的距离之和为dp[root]。

在该题目中，需要求的是各个点到其他点的距离之和。所以在求得其他点到根节点的距离之和后，需要对树进行变化，具体来说便是将根节点下移，子节点上移作为根节点。

假设u为根节点，v为子节点，则有：

代码：

class Solution {
public: 
    
    vector<vector<int>> map1;
    vector<int> dp;
    vector<int> nodes;
    vector<int> ans;
    void dfs(int u, int fa)
    {
        dp[u] = 0;
        nodes[u] = 1;
        for(auto& ch : map1[u])
        {
            if(ch == fa) continue;

            dfs(ch, u);

            dp[u] += dp[ch] + nodes[ch];
            nodes[u] += nodes[ch];
        }
    }

    void dfs2(int u, int fa)
    {
        ans[u] = (dp[u]);
        for(auto& ch : map1[u])
        {
            if(ch == fa) continue;
            int su = dp[u], sc = nodes[u];
            int cu = dp[ch], cc = nodes[ch];
            
            dp[u] -= (dp[ch] + nodes[ch]);
            nodes[u] -= nodes[ch];

            dp[ch] += (dp[u] + nodes[u]);
            nodes[ch] += nodes[u];

            dfs2(ch, u);

            dp[u] = su; nodes[u] = sc;
            dp[ch] = cu; nodes[ch] = cc;
        }
    }
    
    vector<int> sumOfDistancesInTree(int N, vector<vector<int>>& edges) {
        
        dp.resize(N);
        ans.resize(N);
        nodes.resize(N);
        map1.resize(N, {});
        for(auto& edge : edges)
        {
            int u = edge[0], v = edge[1];
            map1[u].emplace_back(v);
            map1[v].emplace_back(u);
        }
        dfs(0, -1);
        dfs2(0, -1);
        return ans;
    }
};