3249. Count the Number of Good Nodes

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.

class Solution {
    int dfs(vector<vector<int>>& adj, int u, int par, int& res) {
        unordered_set<int> us;
        int sum = 1;
        for(auto& v : adj[u]) {
            if(v == par) continue;
            int sub = dfs(adj,v,u,res);
            us.insert(sub);
            sum += sub;
        }
        res += (us.size() <= 1);
        return sum;
    }
public:
    int countGoodNodes(vector<vector<int>>& edges) {
        vector<vector<int>> adj(edges.size() + 1);
        for(auto& e : edges) {
            int u = e[0], v = e[1];
            adj[u].push_back(v);
            adj[v].push_back(u);
        }
        int res = 0;
        dfs(adj,0,-1,res);
        return res;
    }
};