一、Problem
给定一个无向、连通的树。树中有 N 个标记为 0…N-1 的节点以及 N-1 条边 。
第 i 条边连接节点 edges[i][0] 和 edges[i][1] 。
返回一个表示节点 i 与其他所有节点距离之和的列表 ans。
输入: N = 6, edges = [[0,1],[0,2],[2,3],[2,4],[2,5]]
输出: [8,12,6,10,10,10]
解释:
如下为给定的树的示意图：
0
/ \
1 2
/|\
3 4 5
我们可以计算出 dist(0,1) + dist(0,2) + dist(0,3) + dist(0,4) + dist(0,5)
也就是 1 + 1 + 2 + 2 + 2 = 8。 因此，answer[0] = 8，以此类推。
说明: 1 <= N <= 10000
二、Solution
方法一：2 × dfs + 公式推导
最暴力的方法就是枚举所有 x，求递归地出它到其它结点的距离（比如 x->y->z，要求 x 就必须求 y，求 y 就必须求 z），然后将 z、y 的结果返回给 x；对于 y 也是如此，时间复杂度为 O ( n 2 ) O(n^2) O(n2)
思路
设现在 A—B 这条边被断开后：
-
d
i
s
t
(
A
)
dist(A)
dist(A)：表示结点 A 到其子节点
(
y
1
.
.
.
y
n
)
(y_1...y_n)
(y1...yn) 的距离总和
d i s t ( B ) dist(B) dist(B)：同理 -
c
n
t
(
A
)
cnt(A)
cnt(A)：表示结点 A 的子节点数目
c n t ( B ) cnt(B) cnt(B)：同理 -
a
n
s
(
A
)
ans(A)
ans(A)：表示结点 A 到其它结点的距离总和
a n s ( B ) ans(B) ans(B)：同理
不难发现：
{
a
n
s
(
A
)
=
d
i
s
t
(
A
)
+
d
i
s
t
(
B
)
+
c
n
t
(
B
)
①
a
n
s
(
B
)
=
d
i
s
t
(
B
)
+
d
i
s
t
(
A
)
+
c
n
t
(
A
)
②
\left\{ \begin{aligned} ans(A) = dist(A) + dist(B) + cnt(B)\ ① \\ ans(B) = dist(B) + dist(A) + cnt(A)\ ②\\ \end{aligned} \right.
{ans(A)=dist(A)+dist(B)+cnt(B) ①ans(B)=dist(B)+dist(A)+cnt(A) ②
附：（加上一个 cnt [B]是因为 A 到 B 的每一个子节点的距离都是在 B 到子节点的距离基础上 +1，总共有 cnt[B] 个子节点）
一式减二式得
{
a
n
s
(
A
)
=
a
n
s
(
B
)
+
c
n
t
(
B
)
−
c
n
t
(
A
)
③
=
a
n
s
(
B
)
+
(
N
−
c
n
t
(
A
)
)
−
c
n
t
(
A
)
④
\left\{ \begin{aligned} ans(A) = ans(B) + cnt(B) - cnt(A)\ ③\\ =ans(B)+(N-cnt(A))-cnt(A)\ ④ \end{aligned} \right.
{ans(A)=ans(B)+cnt(B)−cnt(A) ③=ans(B)+(N−cnt(A))−cnt(A) ④
对于某个父亲节点 fa 与其孩子结点 child 也可推出
a
n
s
(
c
h
i
l
d
)
=
a
n
s
(
f
a
)
+
(
N
−
c
n
t
(
c
h
i
l
d
)
)
−
c
n
t
(
c
h
i
l
d
)
⑤
ans(child)=ans(fa) + (N-cnt(child))-cnt(child)\ ⑤
ans(child)=ans(fa)+(N−cnt(child))−cnt(child) ⑤
算法
这里需要两步计算：
- 第一步：根据 ①、② 式与子节点 child 与父节点 fa 算出的数组 ans、cnt
- 第二步：根据 ⑤ 式子算出最终的 ans
class Solution {
public:
int n;
vector<vector<int>> g;
vector<int> ans, cnt;
void dfs(int u, int fa) {
for (int& v : g[u]) if (v != fa) {
dfs(v, u);
cnt[u] = cnt[u] + cnt[v];
ans[u] = ans[u] + ans[v] + cnt[v]; //结点u的距离=所有子节点距离和+子节点数量
}
}
void dfs1(int u, int fa) {
for (int& v : g[u]) if (v != fa) {
ans[v] = ans[u] + n-cnt[v] - cnt[v];
dfs1(v, u);
}
}
vector<int> sumOfDistancesInTree(int N, vector<vector<int>>& es) {
g.resize(N), ans.resize(N), cnt.resize(N, 1);
this->n = N;
for (auto& e : es) {
g[e[0]].push_back(e[1]);
g[e[1]].push_back(e[0]);
}
dfs(0, -1);
dfs1(0, -1);
return ans;
}
};
复杂度分析
- 时间复杂度： O ( n ) O(n) O(n)，
- 空间复杂度： O ( n ) O(n) O(n)
212
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
