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

Question

Given a tree (i.e. a connected, undirected graph that has no cycles) consisting of n nodes numbered from 0 to n - 1 and exactly n - 1 edges. The root of the tree is the node 0, and each node of the tree has a label which is a lower-case character given in the string labels (i.e. The node with the number i has the label labels[i]).

The edges array is given on the form edges[i] = [ai, bi], which means there is an edge between nodes ai and bi in the tree.

Return an array of size n where ans[i] is the number of nodes in the subtree of the ith node which have the same label as node i.

A subtree of a tree T is the tree consisting of a node in T and all of its descendant nodes.

LeetCodeleetcode.com/problems/number-of-nodes-in-the-sub-tree-with-the-same-label

Solution

/*
Algorithm: 类似于tree post-order traversal。
把每一个子节点的int[]加起来计算父节点int[];
一次dfs, Time: O(n), Space: O(n)
*/
int[] ret;
public int[] countSubTrees(int n, int[][] edges, String labels) {
    Map<Integer, List<Integer>> nbr = new HashMap<>();
    for (int[] e : edges) {
        nbr.computeIfAbsent(e[0], o->new ArrayList<>()).add(e[1]);
        nbr.computeIfAbsent(e[1], o->new ArrayList<>()).add(e[0]);
    }        
    ret = new int[n];
    dfs(nbr, -1, 0, labels.toCharArray(), n);
    return ret;            
}

private int[] dfs(Map<Integer, List<Integer>> nbr, int pnt, int cnt, char[] c, int n) {
    int[] num = new int[26];
    num[c[cnt]-'a']++;
    for (int nxt : nbr.get(cnt)) {
        if (nxt == pnt) continue;
        int[] tmp = dfs(nbr, cnt, nxt, c, n);
        for (int i = 0; i < 26; i++) num[i] += tmp[i];
    }
    ret[cnt] = num[c[cnt]-'a'];
    return num;
}