这道题我自己的想法只有对每个点都用一遍Dijkstra然后再求和,显然会超时,所以我都没有尝试。
研究了一下题解,发现题解很巧妙,自己对树的处理还是太稚嫩,之前树链剖分学的都忘光了。
对于固定根节点的,我们应该使用树状dp:
dp[u]=∑v∈son(u)dp[v]+sz[v]dp[u]=\sum_{v\in son(u)} dp[v]+sz[v] dp[u]=v∈son(u)∑dp[v]+sz[v]
其中,dp[u]dp[u]dp[u]表示以u为根节点的子树中根节点u到其儿子节点的所有距离之和,之所以加上sz[v]sz[v]sz[v]是因为长度为1,如果长度不固定的话只需要乘上长度就可以了。
通过上面的dp我们可以算出以一个节点为根节点的答案,但是其他节点呢?朴素的想法是对其他节点每个也进行相同的操作,这样的时间复杂度是O(n2)O(n^2)O(n2),好像还可以接受。
题解给出了更优秀的做法:假设我们已经求出了以节点u为根节点的树的答案dp[u]dp[u]dp[u],对于其直接孩子v,我们是有办法直接求出以v为根节点的答案的。
最主要的性质是:如果更换了根节点,只会影响这两个相邻节点的dp值,对其他节点的dp值是不会有影响的。因此我们只需要更新这两个节点即可。
首先我们应该从dp[u]dp[u]dp[u]中减去dp[v]+sz[v]dp[v]+sz[v]dp[v]+sz[v],因为这个时候是以v为根节点的,然后更新sz[u]sz[u]sz[u],然后再给dp[v]dp[v]dp[v]加上dp[u]+sz[u]dp[u]+sz[u]dp[u]+sz[u],再更新sz[v]sz[v]sz[v]。这样就得到了正确的答案。
代码实现的话首先进行一次树状dp,然后再进行回溯。
发现题解中的代码很精炼,仔细学习了一下自己实现了一个差不多的,几乎没有区别。学到了使用emplace_backemplace\_backemplace_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;}
};