LeetCode 834.树中距离之和 Sum of distance in tree
给定一个无向、连通的树。树中有 N 个标记为 0...N-1 的节点以及 N-1 条边 。
第 i 条边连接节点 edges[i][0] 和 edges[i][1] 。
返回一个表示节点 i 与其他所有节点距离之和的列表 ans。
示例 1:
输入: 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
题目虽然是hard级的,但是感觉难度其实不高。属于较为基础的树状
D
P
DP
DP
首先题目是计算以各点开始的所有距离之和。如果直接计算两点距离。最快的方案是弗洛伊德算法,这是
O
(
N
3
)
O(N^3)
O(N3)的,直接挂
下面分析一下思路:
假设某个节点
r
o
o
t
root
root,令
d
p
[
r
o
o
t
]
[
0
]
dp[root][0]
dp[root][0]代表以当前节点root为根节点,他下面的子树到他的距离和。
d
p
[
r
o
o
t
]
[
1
]
dp[root][1]
dp[root][1]代表以当前节点
r
o
o
t
root
root为根节点的话,他下面的所有节点个数(包含他自己),对于叶子节点。
d
p
[
r
o
o
t
]
[
0
]
=
0
,
d
p
[
r
o
o
t
]
[
1
]
=
1
dp[root][0]=0,dp[root][1]=1
dp[root][0]=0,dp[root][1]=1
另外为了方便处理,推荐固定以0号节点为全局根节点,因为题目保证了树的连通性,所以以哪个点做根节点其实无所谓。而且题目保证了是树,所以不会出现某个节点有多个父节点的可能
不难看出
d
p
[
r
o
o
t
]
[
0
]
=
∑
l
e
a
f
∈
c
h
i
l
d
r
e
n
(
d
p
[
l
e
a
f
]
[
0
]
+
d
p
[
l
e
a
f
]
[
1
]
)
dp[root][0]=\sum_{leaf∈children}{(dp[leaf][0]+dp[leaf][1])}
dp[root][0]=∑leaf∈children(dp[leaf][0]+dp[leaf][1])
d
p
[
r
o
o
t
]
[
1
]
=
(
∑
l
e
a
f
∈
c
h
i
l
d
r
e
n
d
p
[
l
e
a
f
]
[
1
]
)
+
1
dp[root][1]=(\sum_{leaf∈children}{dp[leaf][1]})+1
dp[root][1]=(∑leaf∈childrendp[leaf][1])+1
看不懂上面两个式子可以在纸上画棵树推导一下
此时
d
p
[
0
]
[
0
]
dp[0][0]
dp[0][0]的结果就是0号节点的答案,可以记做
a
n
s
[
0
]
ans[0]
ans[0]。假设
c
h
i
l
d
x
child_x
childx是某个和0相连的节点。如果以
c
h
i
l
d
x
child_x
childx为根节点计算距离的话,因为他和0号节点距离为1,那么不难看出,
c
h
i
l
d
x
child_x
childx自己所覆盖的所有子节点距离
c
h
i
l
d
x
child_x
childx的距离都会少1,而其它的节点距离会增加1,那么刚刚已经计算出
d
p
[
c
h
i
l
d
x
]
[
1
]
dp[child_x][1]
dp[childx][1]其实就是自己的所有子节点数量,所以剩余的节点数量为
N
−
d
p
[
c
h
i
l
d
x
]
[
1
]
N-dp[child_x][1]
N−dp[childx][1]
可以得到方程:
a
n
s
[
c
h
i
l
d
x
]
=
a
n
s
[
0
]
−
d
p
[
c
h
i
l
d
x
]
[
1
]
+
(
N
−
d
p
[
c
h
i
l
d
x
]
[
1
]
)
ans[child_x]=ans[0]-dp[child_x][1]+(N-dp[child_x][1])
ans[childx]=ans[0]−dp[childx][1]+(N−dp[childx][1])
化简一下:
a
n
s
[
c
h
i
l
d
x
]
=
a
n
s
[
0
]
−
2
∗
d
p
[
c
h
i
l
d
x
]
[
1
]
+
N
ans[child_x]=ans[0]-2*dp[child_x][1]+N
ans[childx]=ans[0]−2∗dp[childx][1]+N
这里只是关于0号节点的转移方程,因为最初只计算了
a
n
s
[
0
]
ans[0]
ans[0],有了
a
n
s
[
0
]
ans[0]
ans[0]根据上面的式子可以计算与0相连的所有儿子节点的
a
n
s
[
c
h
i
l
d
x
]
ans[child_x]
ans[childx]计算出所有
a
n
s
[
c
h
i
l
d
x
]
ans[child_x]
ans[childx]后,可以进一步计算孙子节点从而计算整棵树
一般的动态转移方程:
a
n
s
[
l
e
a
f
i
]
=
a
n
s
[
r
o
o
t
]
−
2
∗
d
p
[
l
e
a
f
i
]
[
1
]
+
N
ans[leaf_i]=ans[root]-2*dp[leaf_i][1]+N
ans[leafi]=ans[root]−2∗dp[leafi][1]+N
解法的话,可以先用个
d
f
s
dfs
dfs计算出
d
p
[
i
]
[
j
]
dp[i][j]
dp[i][j],然后再次使用
d
f
s
dfs
dfs,计算所有
a
n
s
ans
ans,两次
d
f
s
dfs
dfs都是
O
(
N
)
O(N)
O(N)的
我的代码,80ms:
class Solution{
private:
vector<int> graph[10005];
int dp[10005][2];
int ans[10005];
public:
void init_graph(int N,vector<vector<int>>& edges)
{
for(int i=0;i<N;++i)
{
graph[i].clear();
dp[i][0]=0;dp[i][1]=1;
}
for(int i=0;i<edges.size();++i)
{
graph[edges[i][0]].push_back(edges[i][1]);
graph[edges[i][1]].push_back(edges[i][0]);
}
}
pair<int,int> dfs_fill_dp(int root,int father)
{
int weight_sum=0,node_sum=1;
for(int i=0;i<graph[root].size();++i)
{
if(graph[root][i]==father)
continue;
pair<int,int> result=dfs_fill_dp(graph[root][i],root);
weight_sum=weight_sum+result.first+result.second;
node_sum=node_sum+result.second;
}
dp[root][0]=weight_sum,dp[root][1]=node_sum;
return make_pair(dp[root][0],dp[root][1]);
}
void fill_ans_array(int root,int father,int N)
{
for(int i=0;i<graph[root].size();++i)
{
if(graph[root][i]==father)
continue;
ans[graph[root][i]]=ans[root]+N-2*dp[graph[root][i]][1];
fill_ans_array(graph[root][i],root,N);
}
}
vector<int> sumOfDistancesInTree(int N,vector<vector<int>>& edges)
{
vector<int> result;
init_graph(N,edges);
pair<int,int> root_result=dfs_fill_dp(0,-1);
ans[0]=root_result.first;
fill_ans_array(0,-1,N);
for(int i=0;i<N;++i)
result.push_back(ans[i]);
return result;
}
};