Problem
An undirected, connected tree with N nodes labelled 0…N-1 and N-1 edges are given.
The ith edge connects nodes edges[i][0] and edges[i][1] together.
Return a list ans, where ans[i] is the sum of the distances between node i and all other nodes.
Note: 1 <= N <= 10000
Example
Solution
class Solution {
public:
unordered_map
<int, vector
<int>> graph
;
unordered_map
<int, int> chs
;
void help(int N
,int p
) {
int r
= 0;
for (auto& e
: graph
[N
]) {
if (p
== e
)
continue;
help(e
, N
);
r
+= (1 + chs
[e
]);
}
chs
[N
] = r
;
return ;
}
int fun(int N
, int p
) {
int sum
= 0;
for (auto& e
: graph
[N
]) {
if (p
== e
)
continue;
sum
+=(fun(e
, N
)+chs
[e
]);
sum
++;
}
return sum
;
}
vector
<int> ret
;
int totalNodes
= 0;
void dfs(int N
,int p
) {
for (auto& e
: graph
[N
]) {
if (p
!=-1) {
ret
[N
] = ( ret
[p
] - chs
[N
] + (totalNodes
-(chs
[N
]+1+1)));
}
if (p
== e
)
continue;
dfs(e
, N
);
}
}
vector
<int> sumOfDistancesInTree(int N
, vector
<vector
<int>>& edges
) {
if (N
== 1) {
return { 0 };
}
ret
= vector
<int>(N
, 0);
totalNodes
= N
;
for (auto& e
: edges
) {
graph
[e
[0]].push_back(e
[1]);
graph
[e
[1]].push_back(e
[0]);
}
help(0,-1);
int disRoot
= fun(0, -1);
ret
[0] = disRoot
;
dfs(0, -1);
return ret
;
}
};
