Count the Number of Good Nodes

medium tree dfs

Problem

There is an undirected tree with n nodes labeled 0 to n-1, rooted at node 0, given by an edge list. A node is good if all of the subtrees rooted at its children have the same size. Return the number of good nodes. (A leaf has no children, so it is trivially good.)

Inputedges = [[0,1],[0,2],[1,3],[1,4],[2,5],[2,6]]

Output7

Every node here has children whose subtrees are equal in size (or has 0/1 child), so all 7 nodes are good.

edges as [u,v] list

def count_good_nodes(edges):
    n = len(edges) + 1
    g = [[] for _ in range(n)]
    for a, b in edges:
        g[a].append(b); g[b].append(a)
    good = 0
    def dfs(node, parent):
        nonlocal good
        total, child_sizes = 1, []
        for nx in g[node]:
            if nx != parent:
                s = dfs(nx, node)
                child_sizes.append(s)
                total += s
        if all(c == child_sizes[0] for c in child_sizes):
            good += 1
        return total
    dfs(0, -1)
    return good

function countGoodNodes(edges) {
  const n = edges.length + 1;
  const g = Array.from({length: n}, () => []);
  for (const [a, b] of edges) { g[a].push(b); g[b].push(a); }
  let good = 0;
  function dfs(node, parent) {
    let total = 1; const childSizes = [];
    for (const nx of g[node]) if (nx !== parent) {
      const s = dfs(nx, node);
      childSizes.push(s);
      total += s;
    }
    if (childSizes.every(c => c === childSizes[0])) good++;
    return total;
  }
  dfs(0, -1);
  return good;
}

import java.util.*;
class Solution {
    List<Integer>[] g; int good = 0;
    public int countGoodNodes(int[][] edges) {
        int n = edges.length + 1;
        g = new List[n];
        for (int i = 0; i < n; i++) g[i] = new ArrayList<>();
        for (int[] e : edges) { g[e[0]].add(e[1]); g[e[1]].add(e[0]); }
        dfs(0, -1);
        return good;
    }
    int dfs(int node, int parent) {
        int total = 1; List<Integer> childSizes = new ArrayList<>();
        for (int nx : g[node]) if (nx != parent) {
            int s = dfs(nx, node);
            childSizes.add(s); total += s;
        }
        boolean ok = true;
        for (int c : childSizes) if (c != childSizes.get(0)) ok = false;
        if (ok) good++;
        return total;
    }
}

class Solution {
public:
    vector<vector<int>> g; int good = 0;
    int dfs(int node, int parent) {
        int total = 1; vector<int> childSizes;
        for (int nx : g[node]) if (nx != parent) {
            int s = dfs(nx, node);
            childSizes.push_back(s); total += s;
        }
        bool ok = true;
        for (int c : childSizes) if (c != childSizes[0]) ok = false;
        if (ok) good++;
        return total;
    }
    int countGoodNodes(vector<vector<int>>& edges) {
        int n = edges.size() + 1;
        g.assign(n, {});
        for (auto& e : edges) { g[e[0]].push_back(e[1]); g[e[1]].push_back(e[0]); }
        dfs(0, -1);
        return good;
    }
};

Explanation

"Good" depends only on the sizes of a node's child subtrees: all of them must be equal. So if we know every subtree size, checking goodness is a quick comparison.

We root the tree at node 0 and run a single post-order DFS. Each call returns the size of the subtree it roots. For a node, we recurse into every child (every neighbor that isn't the parent), collecting their returned sizes and summing them plus one for the node itself.

Once we have the list of child sizes, the node is good when they are all equal to the first one. Nodes with zero or one child pass this test automatically, which correctly makes every leaf good.

We increment a shared counter whenever a node turns out good, and return the subtree total so the parent can use it.

Example: in the sample tree, node 0 has two children whose subtrees both have size 3, node 1 and node 2 each have two size-1 children, and the leaves have none. Every node is good, so the answer is 7.

Time: O(n) Space: O(n)