leetcode 834. 树中距离之和

题目详情

题目链接 给定一个无向、连通的树。树中有 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

我的代码

class Solution { public: vector<int> sumOfDistancesInTree(int N, vector<vector<int>>& edges) { if (N == 1) { return {0}; } if (N == 2) { return {1, 1}; } vector<vector<int>> dist(N, vector<int> (N, N)); for (auto & edge : edges) { dist[edge[0]][edge[1]] = 1; dist[edge[1]][edge[0]] = 1; } for (int i = 0; i < N; ++i) { dist[i][i] = 0; } for (int m = 0; m < N; ++m) { for (int i = 0; i < N; ++i) { for (int j = 0; j < N; ++j) { dist[i][j] = min(dist[i][j], dist[i][m] + dist[m][j]); dist[j][i] = dist[i][j]; } } } vector<int> result(N, 0); for (int i = 0; i < N; ++i) { for (int j = 0; j < N; ++j) { result[i] += dist[i][j]; } } return result; } };

我的成绩

执行结果：超出时间限制 64 / 69 个通过测试用例

一些想法

我使用的弗洛伊德算法，但是过不去。。。

执行用时为 84 ms 的范例

class Solution { public: int num[10010],cost[10010]; vector<int> res; int e[10010],ne[20010],to[20010],idx; void dfs(int n,int f) { for(int i=e[n];i;i=ne[i]) { if(to[i]==f) continue; dfs(to[i],n); num[n]+=num[to[i]]; cost[n]+=cost[to[i]]+num[to[i]]; } num[n]++; } void dfs1(int n,int f,int c) { res[n]=c; for(int i=e[n];i;i=ne[i]) { if(to[i]==f) continue; dfs1(to[i],n,c+num[0]-num[to[i]]-num[to[i]]); } } vector<int> sumOfDistancesInTree(int N, vector<vector<int>>& edges) { res.resize(N); for(auto &c:edges) { to[++idx]=c[1]; ne[idx]=e[c[0]]; e[c[0]]=idx; to[++idx]=c[0]; ne[idx]=e[c[1]]; e[c[1]]=idx; } dfs(0,-1); dfs1(0,-1,cost[0]); return res; } };

思考

参见官方解答