这道题我自己的想法只有对每个点都用一遍Dijkstra然后再求和,显然会超时,所以我都没有尝试。
研究了一下题解,发现题解很巧妙,自己对树的处理还是太稚嫩,之前树链剖分学的都忘光了。
对于固定根节点的,我们应该使用树状dp:
d
p
[
u
]
=
∑
v
∈
s
o
n
(
u
)
d
p
[
v
]
+
s
z
[
v
]
dp[u]=\sum_{v\in son(u)} dp[v]+sz[v]
dp[u]=v∈son(u)∑dp[v]+sz[v]
其中,
d
p
[
u
]
dp[u]
dp[u]表示以u为根节点的子树中根节点u到其儿子节点的所有距离之和,之所以加上
s
z
[
v
]
sz[v]
sz[v]是因为长度为1,如果长度不固定的话只需要乘上长度就可以了。
通过上面的dp我们可以算出以一个节点为根节点的答案,但是其他节点呢?朴素的想法是对其他节点每个也进行相同的操作,这样的时间复杂度是 O ( n 2 ) O(n^2) O(n2),好像还可以接受。
题解给出了更优秀的做法:假设我们已经求出了以节点u为根节点的树的答案 d p [ u ] dp[u] dp[u],对于其直接孩子v,我们是有办法直接求出以v为根节点的答案的。
最主要的性质是:如果更换了根节点,只会影响这两个相邻节点的dp值,对其他节点的dp值是不会有影响的。因此我们只需要更新这两个节点即可。
首先我们应该从 d p [ u ] dp[u] dp[u]中减去 d p [ v ] + s z [ v ] dp[v]+sz[v] dp[v]+sz[v],因为这个时候是以v为根节点的,然后更新 s z [ u ] sz[u] sz[u],然后再给 d p [ v ] dp[v] dp[v]加上 d p [ u ] + s z [ u ] dp[u]+sz[u] dp[u]+sz[u],再更新 s z [ v ] sz[v] sz[v]。这样就得到了正确的答案。
代码实现的话首先进行一次树状dp,然后再进行回溯。
发现题解中的代码很精炼,仔细学习了一下自己实现了一个差不多的,几乎没有区别。学到了使用 e m p l a c e _ b a c k emplace\_back emplace_back的复杂度好像更加优秀。
还是要多做hard题目,对自己的提升比较大。
class Solution {
public:
vector<int> dp,sz,ans;
vector< vector<int> > graph;
void Init(int n, vector<vector<int> >&edges)
{
dp.resize(n, 0);
sz.resize(n, 0);
ans.resize(n, 0);
graph.resize(n, {});
int u,v;
for(auto& edge:edges)
{
u = edge[0]; v = edge[1];
graph[u].emplace_back(v);
graph[v].emplace_back(u);
}
}
void dfs(int u,int father)
{
sz[u] = 1; dp[u] = 0;
for(auto& v:graph[u])
{
if(v == father) continue;
dfs(v, u);
sz[u] += sz[v];
dp[u] += dp[v] + sz[v];
}
}
void dfs2(int u, int father)
{
ans[u] = dp[u];
int dpu, dpv, szu, szv;
for(auto& v:graph[u])
{
if(v == father) continue;
dpu = dp[u]; dpv = dp[v];
szu = sz[u]; szv = sz[v];
dp[u] -= dp[v] + sz[v];
sz[u] -= sz[v];
dp[v] += dp[u] + sz[u];
sz[v] += sz[u];
dfs2(v, u);
dp[u] = dpu; dp[v] = dpv;
sz[u] = szu; sz[v] = szv;
}
}
vector<int> sumOfDistancesInTree(int N, vector<vector<int>>& edges) {
Init(N, edges);
dfs(0, -1);
dfs2(0, -1);
return ans;
}
};