834. Sum of Distances in Tree

An undirected, connected tree with N nodes labelled 0...N-1 and N-1 edges are given.

The ith edge connects nodes edges[i][0] and edges[i][1] together.

Return a list ans, where ans[i] is the sum of the distances between node i and all other nodes.

Example 1:

Input: N = 6, edges = [[0,1],[0,2],[2,3],[2,4],[2,5]]
Output: [8,12,6,10,10,10]
Explanation: 
Here is a diagram of the given tree:
  0
 / \
1   2
   /|\
  3 4 5
We can see that dist(0,1) + dist(0,2) + dist(0,3) + dist(0,4) + dist(0,5)
equals 1 + 1 + 2 + 2 + 2 = 8.  Hence, answer[0] = 8, and so on.

Note: 1 <= N <= 10000

題目的input給的是graph中的edges，所以建立一個adjacency list去儲存每個node對應到的所有nodes。可以利用2D vector去完成。 建立兩個vector，count是記錄以每個node的subtree裡有多少個node，res是記錄每個node的distance sum。 第一個pre-order DFS有兩個關係式: 1. count[root] = sum(count[i]) + 1 2. res[root] = sum(count[i]) + sum(res[i]) 可以從觀察簡單的例子得來。 一開始第一個DFS是pre-order，會先從下而上的紀錄資訊，所以這時的res紀錄的是所有的children到該node的distance sum。這時只有res[root]的值是正確的。 再來要從tree的graph由上而下的去更新res vector的值，因為原本只有紀錄由下而上到該node的距離和。現在必須要加上該node往上的所有nodes到該節點的距離和。 從root往下看，i為root連結到的任一node，我們已知res[root]的值是正確的，現在想要更新res[i]，從以i為subtree的所有nodes到i的距離比到root減少1，而i往上的其他剩下的節點到i的距離比到root還增加1。 所以可以得到: res[i] = res[root] - count[i] + (N - count[i])

// Bottom-Up DFS Recursive Helper Function
void postOrder(vector<vector<int> >& tree, int cur, int pre, vector<int>& count, vector<int>& res) {
    for (int i : tree[cur]) {
        if (i == pre) continue;
        postOrder(tree, i, cur, count, res);
        count[cur] += count[i];
        res[cur] += res[i] + count[i];
    }
    ++count[cur]; // add the node itself
}
// Top-Down DFS Recursive Helper Function
void preOrder(vector<vector<int> >& tree, int cur, int pre, vector<int>& count, vector<int>& res) {
    for (int i : tree[cur]) {
        if (i == pre) continue;
        res[i] = res[cur] - count[i] + count.size() - count[i];
        preOrder(tree, i, cur, count, res);
    }
}
vector<int> sumOfDistancesInTree(int N, vector<vector<int>>& edges) { // time: O(N); space: O(N)
    // res: distance for each node
    // count[i]: record how many nodes for a subtree rooted at i
    vector<int> res(N), count(N);
    vector<vector<int> > tree(N); // adjacency list
    // Build the graph
    for (const vector<int>& edge : edges) {
        tree[edge[0]].push_back(edge[1]);
        tree[edge[1]].push_back(edge[0]);
    }
    postOrder(tree, 0, -1, count, res);
    preOrder(tree, 0, -1, count, res);
    return res;
}