An undirected, connected tree with N
nodes labelled 0...N-1
and N-1
edges
are given.
The i
th 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.
Example 1:
Input: N = 6, edges = [[0,1],[0,2],[2,3],[2,4],[2,5]]
Output: [8,12,6,10,10,10]
Explanation:
Here is a diagram of the given tree:
0
/ \
1 2
/|\
3 4 5
We can see that dist(0,1) + dist(0,2) + dist(0,3) + dist(0,4) + dist(0,5)
equals 1 + 1 + 2 + 2 + 2 = 8. Hence, answer[0] = 8, and so on.
Note: 1 <= N <= 10000
class Solution {
public:
vector<int> sumOfDistancesInTree(int N, vector<vector<int>>& edges)
{
if(edges.empty()) return {0};
v.resize(N);
ans.resize(N,0);
for(int i=0;i<edges.size();i++)
{
v[edges[i][0]].connect.push_back(edges[i][1]);
v[edges[i][1]].connect.push_back(edges[i][0]);
}
dfs(edges[0][0]);
ans[edges[0][0]] = v[edges[0][0]].sum;
dfsans(N,edges[0][0]);
return ans;
}
struct TreeNode
{
int sum=0;
int num=1;
vector<int> connect;
bool visited=false;
};
vector<TreeNode> v;
vector<int> ans;
void dfs(int val)
{
v[val].visited = true;
for(int i=0;i<v[val].connect.size();i++)
{
if(!v[v[val].connect[i]].visited)
{
dfs(v[val].connect[i]);
v[val].num += v[v[val].connect[i]].num;
v[val].sum += v[v[val].connect[i]].sum + v[v[val].connect[i]].num;
}
}
}
void dfsans(int N,int val)
{
for(int i=0;i<v[val].connect.size();i++)
{
if(ans[v[val].connect[i]] == 0)
{
ans[v[val].connect[i]] = ans[val] + N - 2*v[v[val].connect[i]].num;
dfsans(N,v[val].connect[i]);
}
}
}
};