[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> &graph, vector &child, vector &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> &graph, vector &ans, vector &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 sumOfDistancesInTree(int N, vector>& edges) {
		vector>graph(N);
		for (vector &e : edges) {
			graph[e[0]].push_back(e[1]);
			graph[e[1]].push_back(e[0]);
		}
		vectorchild(N, 1);
		vectorans(N, 0);
		DFS(graph, child,ans, 0, -1);
		dfs(graph, ans, child, N, 0, -1);
        return ans;
	}
};

你可能感兴趣的:(Tree)