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);
}
|