aboutsummaryrefslogtreecommitdiff
path: root/count_num_good_nodes.c
blob: 44879ed1c704edba25e41ef328e402861d21d154 (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
65
66
67
68
69
70
71
72
73
74
// 3249. Count the Number of Good Nodes
#include "leetcode.h"

/*
 * there is an undirected tree with 'n' nodes labeled from 0 to 'n - 1', and
 * rooted at node 0. you are given a 2d integer array 'edges' of length 'n - 1'
 * where 'edges[i] = [a_i, b_i]' indicates that there is an edge between nodes
 * 'a_i' and 'b_i' in the tree. a node is good if all the subtrees rooted at is
 * children have the same size. return the number of good nodes in the given
 * tree. a subtree of 'treeName' is a tree consisting of a node in 'treeName'
 * and all of its descendants.
 */

struct map {
  int *next;
  int cnt;
};

int dfs(struct map *m, int fa, int curr, int *res) {
  int n = m[curr].cnt;
  if (n == 1 && curr) {
    (*res)++;
    return 1;
  }
  int prev = -1, sum = 0;
  bool flag = true;
  for (int i = 0; i < n; i++) {
    int next = m[curr].next[i];
    if (next != fa) {
      int x = dfs(m, curr, next, res);
      if (prev == -1)
        prev = x;
      else if (x != prev)
        flag = false;
      sum += x;
    }
  }
  *res += flag;
  return sum + 1;
}

int countGoodNodes(int **edges, int edgesSize, int *edgesColSize) {
  int n = edgesSize + 1;
  int *cnt_child = (int *)calloc(n, sizeof(int));
  for (int i = 0; i < edgesSize; i++) {
    int n1 = edges[i][0], n2 = edges[i][1];
    cnt_child[n1]++;
    cnt_child[n2]++;
  }
  struct map *m = (struct map *)calloc(n, sizeof(struct map));
  for (int i = 0; i < n; i++)
    m[i].next = (int *)calloc(cnt_child[i], sizeof(int));
  for (int i = 0; i < edgesSize; i++) {
    int n1 = edges[i][0], n2 = edges[i][1];
    m[n1].next[m[n1].cnt++] = n2;
    m[n2].next[m[n2].cnt++] = n1;
  }
  int ans = 0;
  dfs(m, -1, 0, &ans);
  return ans;
}

int main() {
  int e1i[6][2] = {{0, 1}, {0, 2}, {1, 3}, {1, 4}, {2, 5}, {2, 6}};
  int e2i[8][2] = {{0, 1}, {1, 2}, {2, 3}, {3, 4},
                   {0, 5}, {1, 6}, {2, 7}, {3, 8}};
  struct two_d_arr e1, e2;
  two_d_arr_init(&e1, 6, 2, e1i);
  two_d_arr_init(&e2, 8, 2, e2i);
  printf("%d\n", countGoodNodes(e1.arr, e1.row_size, e1.col_size)); // expect: 7
  printf("%d\n", countGoodNodes(e2.arr, e2.row_size, e2.col_size)); // expect: 6
  two_d_arr_free(&e1);
  two_d_arr_free(&e2);
}