aboutsummaryrefslogtreecommitdiff
path: root/sum_dist_tree.c
blob: fbd75bc0bc969eb8ba7fdef2f025e2845b53de16 (plain)
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;
}