834. Sum of Distances in Tree

https://leetcode.com/problems/sum-of-distances-in-tree/

问给定树上每个结点到其它任意结点的距离和。虽然不是二叉树，但依然可以分治看，只是不是看左右结点，而是任意结点x都可以看作它和它的parent y把树分成了两边。让x@x代表x到它下面所有结点距离和，#x表x和它下面所有结点的个数，那么可以得出:

res[x] = x@x + y@y + #y res[y] = y@y + x@x + #x res[x] = res[y] + #y - #x = res[y] + (N - #x) - #x

而res[root]求解是很简单的，一次后序遍历即可。那用上面的关系再做一次先序遍历就可以依次得到res[root->children] => res[root->children->children]...

思路参考了这。

/*
 * @lc app=leetcode id=834 lang=cpp
 *
 * [834] Sum of Distances in Tree
 */

// @lc code=start
class Solution {
public:
    vector<int> count, res;
    vector<unordered_set<int>> tree;
    void dfs(int cur, int pre) {
        for (auto nei : tree[cur]) {
            if (nei == pre) continue;
            dfs(nei, cur);
            count[cur] += count[nei];
            // 每个结点到它下面所有结点的距离和 == sum(child到它下面所有结点和 + child以它下面结点个数)
            res[cur] += res[nei] + count[nei];
        }
    }

    void dfs2(int cur, int pre) {
        for (auto nei : tree[cur]) {
            if (nei == pre) continue;
            // res[x] = x@x + y@y + #y
            // res[y] = y@y + x@x + #x
            // res[x] = res[y] + #y - #x = res[y] + (N - #x) - #x 
            res[nei] = res[cur] + (count.size() - count[nei]) - count[nei];
            dfs2(nei, cur);
        }
    }
    vector<int> sumOfDistancesInTree(int N, vector<vector<int>>& edges) {
         tree.resize(N);
         res.assign(N, 0);
         count.assign(N, 1);
         for (const auto &e : edges) {
             tree[e[0]].insert(e[1]);
             tree[e[1]].insert(e[0]);
         }
         dfs(0, -1);
         dfs2(0, -1);
         return res;
    }
};
// @lc code=end