An undirected, connected tree with N
nodes labelled 0...N-1
and N-1
edges
are given.
The i
th 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);
}
}
};