leetcode: 834. Sum of Distances in Tree O(N)时间复杂度

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.

首先建立一个 struct  名字tree;

sum 表示 以这个节点为根的子树的所有子节点到该节点的距离的和；

num 表示 以这个节点为根的子树的所有节点个数；

vector c存储所有与这个数值相关联的数字

ans[i]表示数字i 离其他节点的距离。

首先我们进行深度优先遍历，每遍历一个节点，把该节点标记为已访问。下次在他的孩子节点的 vector c中遇到该节点，就知道它是父亲，不会重复访问。每个节点的sum  num 可由递推公式推导出来

 

				v[val].num += v[i].num;
				v[val].sum += v[i].sum + v[i].num;

val 是父亲节点，i  是子节点。同样的 ans[i]  也可以由父节点val  推导出来

 

ans[i] = ans[val] + N - 2 * v[i].num;

下面是代码： 复杂度 O（n）

class Solution {
public:
	struct tree {
		int sum;
		int num;
		bool visited;
		vector c;
		tree() :sum(0), num(1),visited(false) {};
	};
	vector ans;
	vector v;
	vector sumOfDistancesInTree(int N, vector>& edges) {
		if (edges.empty()) return { 0 };
		v.resize(N);
		ans.resize(N);
		vector f(N, N);
		int rot = edges[0][0];
		for (auto& e : edges) {
			v[e[0]].c.emplace_back(e[1]);
			v[e[1]].c.emplace_back(e[0]);
		} 
		dfs(rot);
		ans[rot] = v[rot].sum;
		dfsAns(rot, N);
		return ans;
	}
private:
	void dfs(int val) {
		v[val].visited = true;
		for (int& i : v[val].c) 
			if (!v[i].visited) {
				dfs(i);
				v[val].num += v[i].num;
				v[val].sum += v[i].sum + v[i].num;
			}
	}
	void dfsAns(int val,int& N) {
		for (int i : v[val].c) 
			if (ans[i] == 0) {
				ans[i] = ans[val] + N - 2 * v[i].num;
				dfsAns(i, N);
			}
	}
};