aboutsummaryrefslogtreecommitdiff
diff options
context:
space:
mode:
-rw-r--r--count_num_good_nodes.c74
-rw-r--r--count_num_good_nodes.py77
2 files changed, 151 insertions, 0 deletions
diff --git a/count_num_good_nodes.c b/count_num_good_nodes.c
new file mode 100644
index 0000000..44879ed
--- /dev/null
+++ b/count_num_good_nodes.c
@@ -0,0 +1,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);
+}
diff --git a/count_num_good_nodes.py b/count_num_good_nodes.py
new file mode 100644
index 0000000..fa876f9
--- /dev/null
+++ b/count_num_good_nodes.py
@@ -0,0 +1,77 @@
+# 3249. Count the Number of Good Nodes
+from collections import defaultdict
+
+"""
+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.
+"""
+
+
+class Solution(object):
+ def countGoodNodes(self, edges):
+ """
+ :type edges: List[List[int]]
+ :rtype: int
+ """
+ n = len(edges) + 1
+ graph = defaultdict(list)
+ subtree_size = [0] * n
+
+ def dfs(curr, parent):
+ n = 1
+ for i in graph[curr]:
+ if i != parent:
+ n += dfs(i, curr)
+ subtree_size[curr] = n
+ return n
+
+ for u, v in edges:
+ graph[u].append(v)
+ graph[v].append(u)
+ dfs(0, -1)
+ ans = 0
+ for i in range(n):
+ flag, prev = True, -1
+ for j in graph[i]:
+ if subtree_size[j] < subtree_size[i]:
+ if prev == -1:
+ prev = subtree_size[j]
+ elif prev != subtree_size[j]:
+ flag = False
+ break
+ if flag:
+ ans += 1
+ return ans
+
+
+if __name__ == "__main__":
+ obj = Solution()
+ print(obj.countGoodNodes(edges=[[0, 1], [0, 2], [1, 3], [1, 4], [2, 5], [2, 6]]))
+ print(
+ obj.countGoodNodes(
+ edges=[[0, 1], [1, 2], [2, 3], [3, 4], [0, 5], [1, 6], [2, 7], [3, 8]]
+ )
+ )
+ print(
+ obj.countGoodNodes(
+ edges=[
+ [0, 1],
+ [1, 2],
+ [1, 3],
+ [1, 4],
+ [0, 5],
+ [5, 6],
+ [6, 7],
+ [7, 8],
+ [0, 9],
+ [9, 10],
+ [9, 12],
+ [10, 11],
+ ]
+ )
+ )