834. Sum of Distances in Tree
There is an undirected connected tree with n nodes labeled from 0 to n - 1 and n - 1 edges.
You are given the integer n and the array edges where edges[i] = [ai, bi] indicates that there is an edge between nodes ai and bi in the tree.
Return an array answer of length n where answer[i] is the sum of the distances between the ith node in the tree and all other nodes.
- new solution update 2022.03.17
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
| class Solution {
vector<vector<int>> adj;
vector<int> res, count;
void helper1(int node, int parent) {
for(auto& near: adj[node]) {
if(near == parent) continue;
helper1(near, node);
count[node] += count[near];
res[node] += count[near] + res[near];
}
}
void helper2(int node, int parent) {
for(auto& near: adj[node]) {
if(near == parent) continue;
res[near] += res[node] - count[near] - res[near] + count.size() - count[near];
helper2(near,node);
}
}
public:
vector<int> sumOfDistancesInTree(int n, vector<vector<int>>& edges) {
adj = vector<vector<int>>(n);
res = vector<int>(n,0);
count = vector<int>(n,1);
for(auto e : edges) {
adj[e[0]].push_back(e[1]);
adj[e[1]].push_back(e[0]);
}
helper1(0,-1);
helper2(0,-1);
return res;
}
};
|
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
| class Solution {
vector<vector<int>> g;
vector<vector<pair<int,int>>> child;
vector<int> res;
pair<int, int> travel(int n, int exclude) {
int nodeCount = 0;
int distance = 0;
for(auto next : g[n]) {
if(next == exclude) continue;
auto [childNodeCount, childDistance] = travel(next, n);
child[n].push_back({childNodeCount, childDistance});
nodeCount += childNodeCount;
distance += childDistance;
}
nodeCount++;
distance += nodeCount;
return {nodeCount, distance};
}
void helper(int n, pair<int, int> fromParent, int exclude) {
auto [selfDistance, totalNodeCount] = fromParent;
bool excluded = false;
for(auto [nodeCount, distance]: child[n]) {
selfDistance += distance;
totalNodeCount += nodeCount;
}
res[n] = selfDistance;
for(int i = 0; i < g[n].size(); i++) {
if(exclude == g[n][i]) {
excluded = true;
continue;
}
int node = totalNodeCount - child[n][i-excluded].first + 1;
int dis = selfDistance - child[n][i-excluded].second + node;
helper(g[n][i], {dis, node}, n);
}
}
public:
vector<int> sumOfDistancesInTree(int n, vector<vector<int>>& edges) {
g = vector<vector<int>>(n);
child = vector<vector<pair<int,int>>>(n);
res = vector<int>(n);
for(auto e : edges) {
g[e[0]].push_back(e[1]);
g[e[1]].push_back(e[0]);
}
travel(0, -1);
helper(0, {0,0},-1);
return res;
}
};
|
