Number of Nodes in the Sub-Tree With the Same Label

MediumHash TableTreeDepth-First SearchBreadth-First SearchCounting

Solution

type CharCount = { [key: string]: number };
 
export function countSubTrees(n: number, edges: number[][], labels: string): number[] {
  const answer = new Array<number>(n).fill(0);
  const trees = new Array(n).fill(undefined).map(() => new Array<number>());
  for (const [a, b] of edges) {
    trees[a].push(b);
    trees[b].push(a);
  }
 
  const dfs = (currentNode: number, parentNode: number): CharCount => {
    const charCount: CharCount = {};
    charCount[labels[currentNode]] = 1;
 
    for (const childNode of trees[currentNode]) {
      if (childNode === parentNode) {
        continue;
      }
      const childCharCount = dfs(childNode, currentNode);
 
      for (const label in childCharCount) {
        if (charCount[label] === undefined) {
          charCount[label] = 0;
        }
        charCount[label] += childCharCount[label];
      }
    }
 
    answer[currentNode] = charCount[labels[currentNode]];
    return charCount;
  };
 
  dfs(0, -1);
  return answer;
}