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
58
59
60
61
62
63
64
|
// 834. Sum of Distances in Tree
#include "leetcode.h"
/*
* 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 'ans' of length 'n' where 'ans[i]'
* is thesum of the distances between the i'th node in the tree and all other
* nodes.
*/
void helper(int *res, int *cnt, int *parent, int i, int n) {
if (res[parent[i]] == -1)
helper(res, cnt, parent, parent[i], n);
res[i] = res[parent[i]] + n - 2 * cnt[i];
}
void change_root(int *parent, int root) {
if (parent[root] == -1)
return;
parent[parent[root]] = root;
parent[root] = -1;
}
int *sumOfDistancesInTree(int n, int **edges, int edgesSize, int *edgesColSize,
int *returnSize) {
int *ans = (int *)malloc(sizeof(int) * (*returnSize = n));
int parent[n], cnt[n];
memset(parent, 255, sizeof(parent));
memset(ans, 255, sizeof(cnt));
memset(cnt, 0, sizeof(cnt));
for (int i = 0; i < edgesSize; i++) {
if (parent[edges[i][1]] == -1)
parent[edges[i][1]] = edges[i][0];
else if (parent[edges[i][0]] == -1)
parent[edges[i][0]] = edges[i][1];
else {
change_root(parent, edges[i][0]);
parent[edges[i][0]] = edges[i][1];
}
}
int sum = 0, root;
for (int i = 0; i < n; i++) {
if (parent[i] == -1)
root = i;
else {
int j = i;
while (parent[j] != -1) {
cnt[j]++;
sum++;
j = parent[j];
}
}
}
cnt[root] = n;
ans[root] = sum;
for (int i = 0; i < n; i++) {
if (ans[i] != -1)
continue;
helper(ans, cnt, parent, i, n);
}
return ans;
}
|