[WXM] LeetCode 834. Sum of Distances in Tree C++

自己开发的博客网站，欢迎访问https://www.weiboke.online
#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

##Approach

1. 题目大意就是返回数组长度为N，然后第i个元素是这个结点i到其他的结点距离的和。这道题数据量比较大，但是还是忍不住暴力试下，然后我直接对每个结点进行BFS直接求答案，到64样例就超时了，最后看了题解明白了，是要进行推算求解，把问题化小，方程ans[node]=ans[parent]-child[node]+(N-child[node])，其中child[node]表示这个node结点下有多少个结点包括自己。
   ##Code

class Solution {
public:
	void DFS(vector<vector<int>> &graph, vector<int> &child, vector<int> &ans, int s, int p) {
		for (int &v : graph[s]) {
			if (v == p)continue;
			DFS(graph, child, ans, v, s);
			child[s] += child[v];
			ans[s] += ans[v] + child[v];
		}
	}
	void dfs(vector<vector<int>> &graph, vector<int> &ans, vector<int> &child,int &N, int s, int p) {
		for (int &v : graph[s]) {
			if (v == p)continue;
			ans[v] = ans[s] - child[v] + N - child[v];
            dfs(graph, ans, child, N, v, s);
		}
	}
	vector<int> sumOfDistancesInTree(int N, vector<vector<int>>& edges) {
		vector<vector<int>>graph(N);
		for (vector<int> &e : edges) {
			graph[e[0]].push_back(e[1]);
			graph[e[1]].push_back(e[0]);
		}
		vector<int>child(N, 1);
		vector<int>ans(N, 0);
		DFS(graph, child,ans, 0, -1);
		dfs(graph, ans, child, N, 0, -1);
        return ans;
	}
};