给定一个无向、连通的树。树中有 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
来源:力扣(LeetCode)
链接:https://leetcode-cn.com/problems/sum-of-distances-in-tree
著作权归领扣网络所有。商业转载请联系官方授权,非商业转载请注明出处。
树形dp,用ans表示所有点到某点距离之和,sum表示某点为根的子树中的点到某点距离之和,cnt表示某点为根的子树中节点数,对于两个相邻点A,B,cntA + cntB = N,ansA = sumA + sumB + cntB,ansB = sumB + sumA + cntA,则ansA = ansB - cntA + cntB = ansB + N - 2*cntA,那么令B为根节点,A为B的子节点,只要只要了ansB和cntB就知道了ansA,首先进行一遍dfs算出每个点cnt,以及根节点的ans,第二遍dfs可以从根往下一点一点算出其他点的ans。
代码:
class Solution { public: vector<int> v[10000],ans,cnt; bool vis[10000] = {false}; int dfs(int k) { vis[k] = true; for(int i = 0;i < v[k].size();i ++) { if(!vis[v[k][i]]) cnt[k] += dfs(v[k][i]); } ans[0] += cnt[k] - 1; vis[k] = false; return cnt[k]; } void get(int l,int k) { vis[k] = true; if(l != -1) ans[k] += ans[l] - 2 * cnt[k]; for(int i = 0;i < v[k].size();i ++) { if(!vis[v[k][i]]) get(k,v[k][i]); } vis[k] = false; } vector<int> sumOfDistancesInTree(int N, vector<vector<int>>& edges) { ans = vector<int>(N,N); cnt = vector<int>(N,1); ans[0] = 0; for(int i = 0;i < N - 1;i ++) { v[edges[i][0]].push_back(edges[i][1]); v[edges[i][1]].push_back(edges[i][0]); } dfs(0); get(-1,0); return ans; } };