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<int> c;
		tree() :sum(0), num(1),visited(false) {};
	};
	vector<int> ans;
	vector<tree> v;
	vector<int> sumOfDistancesInTree(int N, vector<vector<int>>& edges) {
		if (edges.empty()) return { 0 };
		v.resize(N);
		ans.resize(N);
		vector<int> 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);
			}
	}
};