leetcode 834 树中距离之和

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

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

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

示例 1:

输入: 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，以此类推。
说明: 1 <= N <= 10000

来源：力扣（LeetCode）
链接：https://leetcode-cn.com/problems/sum-of-distances-in-tree
著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。

树形dp，用ans表示所有点到某点距离之和，sum表示某点为根的子树中的点到某点距离之和，cnt表示某点为根的子树中节点数，对于两个相邻点A,B，cntA + cntB = N，ansA = sumA + sumB + cntB，ansB = sumB + sumA + cntA，则ansA = ansB - cntA + cntB = ansB + N - 2*cntA，那么令B为根节点，A为B的子节点，只要只要了ansB和cntB就知道了ansA，首先进行一遍dfs算出每个点cnt，以及根节点的ans，第二遍dfs可以从根往下一点一点算出其他点的ans。

代码：

class Solution {
public:
    vector<int> v[10000],ans,cnt;
    bool vis[10000] = {false};
    int dfs(int k) {
        vis[k] = true;
        for(int i = 0;i < v[k].size();i ++) {
            if(!vis[v[k][i]]) cnt[k] += dfs(v[k][i]);
        }
        ans[0] += cnt[k] - 1;
        vis[k] = false;
        return cnt[k];
    }
    void get(int l,int k) {
        vis[k] = true;
        if(l != -1) ans[k] += ans[l] - 2 * cnt[k];
        for(int i = 0;i < v[k].size();i ++) {
            if(!vis[v[k][i]]) get(k,v[k][i]);
        }
        vis[k] = false;
    }
    vector<int> sumOfDistancesInTree(int N, vector<vector<int>>& edges) {
        ans = vector<int>(N,N);
        cnt = vector<int>(N,1);
        ans[0] = 0;
        for(int i = 0;i < N - 1;i ++) {
            v[edges[i][0]].push_back(edges[i][1]);
            v[edges[i][1]].push_back(edges[i][0]);
        }
        dfs(0);
        get(-1,0);
        return ans;
    }
};