3249. Count the Number of Good Nodes

Posted on Edited on In

3249. Count the Number of Good Nodes

Description

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] = [ai, bi] indicates that there is an edge between nodes ai and bi in the tree.

A node is good if all the subtrees rooted at its 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.

Example 1:

1
2
3
Input: edges = [[0,1],[0,2],[1,3],[1,4],[2,5],[2,6]]

Output: 7

Explanation:

All of the nodes of the given tree are good.

Example 2:

1
2
3
Input: edges = [[0,1],[1,2],[2,3],[3,4],[0,5],[1,6],[2,7],[3,8]]

Output: 6

Explanation:

There are 6 good nodes in the given tree. They are colored in the image above.

Example 3:

1
2
3
Input: 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]]

Output: 12

Explanation:

All nodes except node 9 are good.

Constraints:

  • 2 <= n <= 10^5
  • edges.length == n - 1
  • edges[i].length == 2
  • 0 <= ai, bi < n
  • The input is generated such that edges represents a valid tree.

Hints/Notes

Solution

Language: C++

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
class Solution {
public:
    int res = 0;
    vector<vector<int>> graph;

    int countGoodNodes(vector<vector<int>>& edges) {
        int n = edges.size() + 1;
        graph.resize(n, vector<int>());
        for (auto& edge : edges) {
            int u = edge[0], v = edge[1];
            graph[u].push_back(v);
            graph[v].push_back(u);
        }
        traverse(0, -1);
        return res;
    }

    int traverse(int root, int origin) {
        int sum = 1, size = 0;
        bool valid = true;
        for (auto& u : graph[root]) {
            if (u == origin) {
                continue;
            }
            int tmp = traverse(u, root);
            if (size && size != tmp) {
                valid = false;
            }
            size = tmp;
            sum += tmp;
        }
        if (valid) {
            res++;
        }
        return sum;
    }
};